diff --git a/scalene/scalene_arguments.py b/scalene/scalene_arguments.py index 361387f9c..516c1b51c 100644 --- a/scalene/scalene_arguments.py +++ b/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) diff --git a/scalene/scalene_profiler.py b/scalene/scalene_profiler.py index 2844c4082..649c2985f 100644 --- a/scalene/scalene_profiler.py +++ b/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) ) diff --git a/scalene/set_nvidia_gpu_modes.py b/scalene/set_nvidia_gpu_modes.py index 912f9b893..2ec252997 100644 --- a/scalene/set_nvidia_gpu_modes.py +++ b/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 diff --git a/tests/test_coverup_53.py b/tests/test_coverup_53.py index 752bf83d6..d93b542e4 100644 --- a/tests/test_coverup_53.py +++ b/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))