Refactored register files logic; added profile-only files.

plasma-umass · Jan 18, 2024 · bd8ea45 · bd8ea45
1 parent 0f22332
commit bd8ea45
Showing 1 changed file with 17 additions and 15 deletions.
diff --git a/scalene/scalene_profiler.py b/scalene/scalene_profiler.py
@@ -401,14 +401,7 @@ def profile(func: Any) -> Any:
         Scalene.__functions_to_profile[func.__code__.co_filename].add(func)
 
         if Scalene.__args.memory:
-            from scalene import pywhere  # type: ignore
-
-            pywhere.register_files_to_profile(
-                list(Scalene.__files_to_profile),
-                Scalene.__program_path,
-                Scalene.__args.profile_all,
-            )
-
+            Scalene.register_files_to_profile()
         return func
 
     @staticmethod
@@ -1511,6 +1504,7 @@ def should_trace(filename: Filename, func: str) -> bool:
                         return False
 
         # Generic handling follows (when no @profile decorator has been used).
+        # TODO [EDB]: add support for this in traceconfig.cpp
         profile_exclude_list = Scalene.__args.profile_exclude.split(",")
         if any(
             prof in filename for prof in profile_exclude_list if prof != ""
@@ -1826,6 +1820,20 @@ def main() -> None:
         Scalene.set_initialized()
         Scalene.run_profiler(args, left)
 
+    @staticmethod
+    def register_files_to_profile() -> None:
+        """Tells the pywhere module, which tracks memory, which files to profile."""
+        from scalene import pywhere  # type: ignore
+
+        profile_only_list = Scalene.__args.profile_only.split(",")
+
+        pywhere.register_files_to_profile(
+            list(Scalene.__files_to_profile) + profile_only_list,
+            Scalene.__program_path,
+            Scalene.__args.profile_all,
+        )
+
+
     @staticmethod
     def run_profiler(
         args: argparse.Namespace, left: List[str], is_jupyter: bool = False
@@ -1950,13 +1958,7 @@ def run_profiler(
                         Scalene.__program_path = program_path
                     # Grab local and global variables.
                     if Scalene.__args.memory:
-                        from scalene import pywhere  # type: ignore
-
-                        pywhere.register_files_to_profile(
-                            list(Scalene.__files_to_profile),
-                            Scalene.__program_path,
-                            Scalene.__args.profile_all,
-                        )
+                        Scalene.register_files_to_profile()
                     import __main__
 
                     the_locals = __main__.__dict__