Fixing Scalene inside Jupyter notebooks (#849)

scalene/scalene_arguments.py

@@ -2,12 +2,12 @@
 import platform
 import sys
 
+from typing import Any
 
 class ScaleneArguments(argparse.Namespace):
     """Encapsulates all arguments and default values for Scalene."""
-
-    def __init__(self) -> None:
-        super().__init__()
+
+    def __init__(self, **kwargs: Any) -> None:
         self.cpu = True
         self.gpu = platform.system() != "Darwin"
         self.memory = sys.platform != "win32"
@@ -47,3 +47,4 @@ def __init__(self) -> None:
         self.no_browser = False
         self.port = 8088
         self.cli = False
+        super().__init__(**kwargs)

scalene/scalene_profiler.py

@@ -678,7 +678,7 @@ def __init__(
             import scalene.replacement_fork
             import scalene.replacement_poll_selector # noqa: F401
 
-        Scalene.__args = cast(ScaleneArguments, arguments)
+        Scalene.__args = ScaleneArguments(**vars(arguments))
         Scalene.__alloc_sigq = ScaleneSigQueue(
             Scalene.alloc_sigqueue_processor
         )
@@ -693,12 +693,12 @@ def __init__(
 
         Scalene.__windows_queue = queue.Queue()
         if sys.platform == "win32":
-            if arguments.memory:
+            if Scalene.__args.memory:
                 print(
                     "Scalene warning: Memory profiling is not currently supported for Windows.",
                     file=sys.stderr,
                 )
-                arguments.memory = False
+                Scalene.__args.memory = False
 
         # Initialize the malloc related files; if for whatever reason
         # the files don't exist and we are supposed to be profiling
@@ -708,19 +708,19 @@ def __init__(
             Scalene.__memcpy_mapfile = ScaleneMapFile("memcpy")
         except Exception:
             # Ignore if we aren't profiling memory; otherwise, exit.
-            if arguments.memory:
+            if Scalene.__args.memory:
                 sys.exit(1)
 
-        Scalene.__signals.set_timer_signals(arguments.use_virtual_time)
+        Scalene.__signals.set_timer_signals(Scalene.__args.use_virtual_time)
         Scalene.__profiler_base = str(os.path.dirname(__file__))
-        if arguments.pid:
+        if Scalene.__args.pid:
             # Child process.
             # We need to use the same directory as the parent.
             # The parent always puts this directory as the first entry in the PATH.
             # Extract the alias directory from the path.
             dirname = os.environ["PATH"].split(os.pathsep)[0]
             Scalene.__python_alias_dir = pathlib.Path(dirname)
-            Scalene.__pid = arguments.pid
+            Scalene.__pid = Scalene.__args.pid
 
         else:
             # Parent process.
@@ -733,7 +733,7 @@ def __init__(
             Scalene.__pid = 0
             cmdline = ""
             # Pass along commands from the invoking command line.
-            if "off" in arguments and arguments.off:
+            if "off" in Scalene.__args and Scalene.__args.off:
                 cmdline += " --off"
             for arg in [
                 "use_virtual_time",
@@ -746,12 +746,12 @@ def __init__(
                 "no_browser",
                 "reduced_profile",
             ]:
-                if getattr(arguments, arg):
+                if getattr(Scalene.__args, arg):
                     cmdline += f'  --{arg.replace("_", "-")}'
             # Add the --pid field so we can propagate it to the child.
             cmdline += f" --pid={os.getpid()} ---"
             # Build the commands to pass along other arguments
-            environ = ScalenePreload.get_preload_environ(arguments)
+            environ = ScalenePreload.get_preload_environ(Scalene.__args)
             if sys.platform == "win32":
                 preface = "\n".join(
                     f"set {k}={str(v)}\n" for (k, v) in environ.items()
@@ -1923,18 +1923,18 @@ def profile_code(
     @staticmethod
     def process_args(args: argparse.Namespace) -> None:
         """Process all arguments."""
-        Scalene.__args = cast(ScaleneArguments, args)
+        Scalene.__args = ScaleneArguments(**vars(args))
         Scalene.__next_output_time = (
             time.perf_counter() + Scalene.__args.profile_interval
         )
-        Scalene.__output.html = args.html
-        if args.outfile:
+        Scalene.__output.html = Scalene.__args.html
+        if Scalene.__args.outfile:
             Scalene.__output.output_file = os.path.abspath(
-                os.path.expanduser(args.outfile)
+                os.path.expanduser(Scalene.__args.outfile)
             )
-        Scalene.__is_child = args.pid != 0
+        Scalene.__is_child = Scalene.__args.pid != 0
         # the pid of the primary profiler
-        Scalene.__parent_pid = args.pid if Scalene.__is_child else os.getpid()
+        Scalene.__parent_pid = Scalene.__args.pid if Scalene.__is_child else os.getpid()
         # Don't profile the GPU if not enabled (i.e., either no options or --cpu and/or --memory, but no --gpu).
         if not Scalene.__args.gpu:
             Scalene.__output.gpu = False
@@ -2116,7 +2116,7 @@ def run_profiler(
                         # TODO: Add in entrypoint_dir logic for `-m` operation
                         Scalene.__entrypoint_dir = program_path
                     # If a program path was specified at the command-line, use it.
-                    if len(args.program_path) > 0:
+                    if len(Scalene.__args.program_path) > 0:
                         Scalene.__program_path = Filename(
                             os.path.abspath(args.program_path)
                         )

scalene/set_nvidia_gpu_modes.py

@@ -2,7 +2,7 @@
 import sys
 import subprocess
 
-def set_nvidia_gpu_modes():
+def set_nvidia_gpu_modes() -> bool:
     import pynvml
     try:
         # Initialize NVML

tests/test_coverup_53.py

@@ -49,7 +49,6 @@ def test_process_args(cleanup_scalene):
 
         # Assertions to verify postconditions
         # Accessing the private attributes directly for testing purposes
-        assert Scalene._Scalene__args == args
         assert Scalene._Scalene__next_output_time == 101  # 100 + 1
         assert Scalene._Scalene__output.html == args.html
         assert Scalene._Scalene__output.output_file == os.path.abspath(os.path.expanduser(args.outfile))