Merge pull request #52 from F-Secure/add_entrypoints_to_init

Add entrypoints to init
WithSecureOpenSource · Dec 13, 2021 · 7982ba2 · 7982ba2
2 parents f54116f + c58abc6
commit 7982ba2
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 37 deletions.
diff --git a/change_analyzer/__init__.py b/change_analyzer/__init__.py
@@ -3,6 +3,8 @@
 import sys
 
 from gym.envs.registration import register
+from .main import run
+from .sequences_diff import compare
 
 
 register(

diff --git a/change_analyzer/main.py b/change_analyzer/main.py
@@ -42,6 +42,30 @@ def reset() -> WebDriver:
     return driver
 
 
+def run(config: str, steps: int = 0, csv_folder:str = ""):
+    CONFIG.read(config)
+
+    env = gym.make(
+        "app-v0" if CONFIG["driver"]["platform"] == "win" else "web-v0",
+        reset_app=reset,
+    )
+    # For some reason registration sometimes doesn't work and line above can fail, direct class creation as in line
+    # below could be used. See https://github.com/openai/gym/blob/master/docs/creating-environments.md for more info    # about registration.
+    # env = AppEnv(reset, {"Help"})
+    sequence_id = uuid.uuid1()
+    report_dir = "recordings/{}".format(datetime.now().strftime("%Y_%m_%d-%H_%M_%S"))
+    # env = EnhancedMonitor(env, report_dir)
+    env = SequenceRecorder(env, report_dir, sequence_id)
+    try:
+        env.reset()
+        if csv_folder:
+            ReplayAgent(env, csv_folder).run()
+        else:
+            RandomAgent(env, int(steps)).run()
+    finally:
+        env.close()
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -60,28 +84,7 @@ def main():
         required=False,
     )
     args = parser.parse_args()
-    CONFIG.read(args.config)
-
-    env = gym.make(
-        "app-v0" if CONFIG["driver"]["platform"] == "win" else "web-v0",
-        reset_app=reset,
-    )
-    # For some reason registration sometimes doesn't work and line above can fail, direct class creation as in line
-    # below could be used. See https://github.com/openai/gym/blob/master/docs/creating-environments.md for more info
-    # about registration.
-    # env = AppEnv(reset, {"Help"})
-    sequence_id = uuid.uuid1()
-    report_dir = "recordings/{}".format(datetime.now().strftime("%Y_%m_%d-%H_%M_%S"))
-    # env = EnhancedMonitor(env, report_dir)
-    env = SequenceRecorder(env, report_dir, sequence_id)
-    try:
-        env.reset()
-        if args.csv_folder:
-            ReplayAgent(env, args.csv_folder).run()
-        else:
-            RandomAgent(env, int(args.steps)).run()
-    finally:
-        env.close()
+    run(args.config, args.steps, args.csv_folder)
 
 
 if __name__ == "__main__":

diff --git a/change_analyzer/sequences_diff.py b/change_analyzer/sequences_diff.py
@@ -332,6 +332,24 @@ def find_last_two_valid_folders() -> Tuple[str, str]:
 
     return csv_file_paths[0], csv_file_paths[1]
 
+
+def compare(sequence1_folder: str = "", sequence2_folder: str = ""):
+    sequence1_file = ""
+    sequence2_file = ""
+    if not sequence1_folder and not sequence2_folder:
+        sequence1_file, sequence2_file = find_last_two_valid_folders()
+
+    if sequence1_folder and not sequence2_folder:
+        sequence1_file = find_csv_file_within_folder(sequence1_folder)
+        sequence2_file, _ = find_last_two_valid_folders()
+
+    if sequence1_folder and sequence2_folder:
+        sequence1_file = find_csv_file_within_folder(sequence1_folder)
+        sequence2_file = find_csv_file_within_folder(sequence2_folder)
+
+    SequencesDiff(sequence1_file, sequence2_file)
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -345,21 +363,7 @@ def main():
         required=False,
     )
     args = parser.parse_args()
-
-    sequence1_file = ""
-    sequence2_file = ""
-    if not args.sequence1_folder and not args.sequence2_folder:
-        sequence1_file, sequence2_file = find_last_two_valid_folders()
-
-    if args.sequence1_folder and not args.sequence2_folder:
-        sequence1_file = find_csv_file_within_folder(args.sequence1_folder)
-        sequence2_file, _ = find_last_two_valid_folders()
-
-    if args.sequence1_folder and args.sequence2_folder:
-        sequence1_file = find_csv_file_within_folder(args.sequence1_folder)
-        sequence2_file = find_csv_file_within_folder(args.sequence2_folder)
-
-    SequencesDiff(sequence1_file, sequence2_file)
+    compare(args.sequence1_folder, args.sequence2_folder)
 
 
 if __name__ == "__main__":