Refactor and cleaning up demonstration file finding logic + associated test

ml-agents/mlagents/trainers/demo_loader.py

@@ -1,4 +1,3 @@
-import pathlib
 import logging
 import os
 from typing import List, Tuple
@@ -83,39 +82,40 @@ def demo_to_buffer(
     return brain_params, demo_buffer
 
 
-@timed
-def load_demonstration(
-    file_path: str
-) -> Tuple[BrainParameters, List[AgentInfoActionPairProto], int]:
-    """
-    Loads and parses a demonstration file.
-    :param file_path: Location of demonstration file (.demo).
-    :return: BrainParameter and list of AgentInfoActionPairProto containing demonstration data.
+def get_demo_files(path: str) -> List[str]:
     """
+    Retrieves the demonstration file(s) from a path.
+    :param path: Path of demonstration file or directory.
+    :return: List of demonstration files
 
-    # First 32 bytes of file dedicated to meta-data.
-    INITIAL_POS = 33
-    file_paths = []
-    if os.path.isdir(file_path):
-        all_files = os.listdir(file_path)
-        for _file in all_files:
-            if _file.endswith(".demo"):
-                file_paths.append(os.path.join(file_path, _file))
-        if not all_files:
+    Raises errors if |path| is invalid.
+    """
+    if os.path.isfile(path):
+        if not path.endswith(".demo"):
+            raise ValueError("The path provided is not a '.demo' file.")
+        return [path]
+    elif os.path.isdir(path):
+        fnames = filter(lambda fname: fname.endswith(".demo"), os.listdir(path))
+        fpaths = list(map(lambda fname: os.path.join(path, fname), fnames))
+        if len(fpaths) == 0:
             raise ValueError("There are no '.demo' files in the provided directory.")
-    elif os.path.isfile(file_path):
-        file_paths.append(file_path)
-        file_extension = pathlib.Path(file_path).suffix
-        if file_extension != ".demo":
-            raise ValueError(
-                "The file is not a '.demo' file. Please provide a file with the "
-                "correct extension."
-            )
+        return fpaths
     else:
         raise FileNotFoundError(
-            "The demonstration file or directory {} does not exist.".format(file_path)
+            f"The demonstration file or directory {path} does not exist."
         )
 
+
+@timed
+def load_demonstration(
+    path: str,
+) -> Tuple[BrainParameters, List[AgentInfoActionPairProto], int]:
+    """
+    Loads and parses a demonstration file or directory.
+    :param path: File or directory.
+    :return: BrainParameter and list of AgentInfoActionPairProto containing demonstration data.
+    """
+    file_paths = get_demo_files(path)
     brain_params = None
     brain_param_proto = None
     info_action_pairs = []
@@ -131,12 +131,13 @@ def load_demonstration(
                     meta_data_proto = DemonstrationMetaProto()
                     meta_data_proto.ParseFromString(data[pos : pos + next_pos])
                     total_expected += meta_data_proto.number_steps
-                    pos = INITIAL_POS
-                if obs_decoded == 1:
+                    # first 32 bytes of file dedicated to metadata
+                    pos = 33
+                elif obs_decoded == 1:
                     brain_param_proto = BrainParametersProto()
                     brain_param_proto.ParseFromString(data[pos : pos + next_pos])
                     pos += next_pos
-                if obs_decoded > 1:
+                else:
                     agent_info_action = AgentInfoActionPairProto()
                     agent_info_action.ParseFromString(data[pos : pos + next_pos])
                     if brain_params is None:
@@ -149,7 +150,5 @@ def load_demonstration(
                     pos += next_pos
                 obs_decoded += 1
     if not brain_params:
-        raise RuntimeError(
-            f"No BrainParameters found in demonstration file at {file_path}."
-        )
+        raise RuntimeError(f"No BrainParameters found in demonstration file at {path}.")
     return brain_params, info_action_pairs, total_expected

ml-agents/mlagents/trainers/tests/test_demo_loader.py

@@ -1,6 +1,11 @@
 import os
+import pytest
 
-from mlagents.trainers.demo_loader import load_demonstration, demo_to_buffer
+from mlagents.trainers.demo_loader import (
+    load_demonstration,
+    demo_to_buffer,
+    get_demo_files,
+)
 
 
 def test_load_demo():
@@ -27,3 +32,11 @@ def test_load_demo_dir():
 
     _, demo_buffer = demo_to_buffer(path_prefix + "/test_demo_dir", 1)
     assert len(demo_buffer["actions"]) == total_expected - 1
+
+
+def test_nonexisting_paths():
+    path_prefix = os.path.dirname(os.path.abspath(__file__))
+    with pytest.raises(FileNotFoundError):
+        get_demo_files(os.path.join(path_prefix, "idontexistfile.demo"))
+    with pytest.raises(FileNotFoundError):
+        get_demo_files(os.path.join(path_prefix, "idontexistdirectory"))