Codex/remove pyi files and add type annotations

line_profiler/__main__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .line_profiler import main
 
 if __name__ == '__main__':

line_profiler/autoprofile/ast_profile_transformer.py

@@ -1,7 +1,12 @@
+from __future__ import annotations
+
 import ast
+from typing import Union
 
 
-def ast_create_profile_node(modname, profiler_name='profile', attr='add_imported_function_or_module'):
+def ast_create_profile_node(
+        modname: str, profiler_name: str = 'profile',
+        attr: str = 'add_imported_function_or_module') -> ast.Expr:
     """Create an abstract syntax tree node that adds an object to the profiler to be profiled.
 
     An abstract syntax tree node is created which calls the attr method from profile and
@@ -45,7 +50,9 @@ class AstProfileTransformer(ast.NodeTransformer):
     immediately after the import.
     """
 
-    def __init__(self, profile_imports=False, profiled_imports=None, profiler_name='profile'):
+    def __init__(self, profile_imports: bool = False,
+                 profiled_imports: list[str] | None = None,
+                 profiler_name: str = 'profile') -> None:
         """Initializes the AST transformer with the profiler name.
 
         Args:
@@ -63,7 +70,9 @@ def __init__(self, profile_imports=False, profiled_imports=None, profiler_name='
         self._profiled_imports = profiled_imports if profiled_imports is not None else []
         self._profiler_name = profiler_name
 
-    def _visit_func_def(self, node):
+    def _visit_func_def(
+            self, node: ast.FunctionDef | ast.AsyncFunctionDef
+    ) -> ast.FunctionDef | ast.AsyncFunctionDef:
         """Decorate functions/methods with profiler.
 
         Checks if the function/method already has a profile_name decorator, if not, it will append
@@ -91,7 +100,10 @@ def _visit_func_def(self, node):
 
     visit_FunctionDef = visit_AsyncFunctionDef = _visit_func_def
 
-    def _visit_import(self, node):
+    def _visit_import(
+            self, node: ast.Import | ast.ImportFrom
+    ) -> (ast.Import | ast.ImportFrom
+          | list[ast.Import | ast.ImportFrom | ast.Expr]):
         """Add a node that profiles an import
 
         If profile_imports is True and the import is not in profiled_imports,
@@ -121,7 +133,9 @@ def _visit_import(self, node):
             visited.append(expr)
         return visited
 
-    def visit_Import(self, node):
+    def visit_Import(
+            self, node: ast.Import
+    ) -> ast.Import | list[ast.Import | ast.Expr]:
         """Add a node that profiles an object imported using the "import foo" sytanx
 
         Args:
@@ -137,7 +151,9 @@ def visit_Import(self, node):
         """
         return self._visit_import(node)
 
-    def visit_ImportFrom(self, node):
+    def visit_ImportFrom(
+            self, node: ast.ImportFrom
+    ) -> ast.ImportFrom | list[ast.ImportFrom | ast.Expr]:
         """Add a node that profiles an object imported using the "from foo import bar" syntax
 
         Args:

line_profiler/autoprofile/ast_tree_profiler.py

@@ -1,5 +1,8 @@
+from __future__ import annotations
+
 import ast
 import os
+from typing import Type
 
 from .ast_profile_transformer import (AstProfileTransformer,
                                       ast_create_profile_node)
@@ -20,11 +23,11 @@ class AstTreeProfiler:
     """
 
     def __init__(self,
-                 script_file,
-                 prof_mod,
-                 profile_imports,
-                 ast_transformer_class_handler=AstProfileTransformer,
-                 profmod_extractor_class_handler=ProfmodExtractor):
+                 script_file: str,
+                 prof_mod: list[str],
+                 profile_imports: bool,
+                 ast_transformer_class_handler: Type = AstProfileTransformer,
+                 profmod_extractor_class_handler: Type = ProfmodExtractor) -> None:
         """Initializes the AST tree profiler instance with the script file path
 
         Args:
@@ -52,7 +55,8 @@ def __init__(self,
         self._profmod_extractor_class_handler = profmod_extractor_class_handler
 
     @staticmethod
-    def _check_profile_full_script(script_file, prof_mod):
+    def _check_profile_full_script(
+            script_file: str, prof_mod: list[str]) -> bool:
         """Check whether whole script should be profiled.
 
         Checks whether path to script has been passed to prof_mod indicating that
@@ -76,7 +80,7 @@ def _check_profile_full_script(script_file, prof_mod):
         return profile_full_script
 
     @staticmethod
-    def _get_script_ast_tree(script_file):
+    def _get_script_ast_tree(script_file: str) -> ast.Module:
         """Generate an abstract syntax from a script file.
 
         Args:
@@ -93,10 +97,10 @@ def _get_script_ast_tree(script_file):
         return tree
 
     def _profile_ast_tree(self,
-                          tree,
-                          tree_imports_to_profile_dict,
-                          profile_full_script=False,
-                          profile_imports=False):
+                          tree: ast.Module,
+                          tree_imports_to_profile_dict: dict[int, str],
+                          profile_full_script: bool = False,
+                          profile_imports: bool = False) -> ast.Module:
         """Add profiling to an abstract syntax tree.
 
         Adds nodes to the AST that adds the specified objects to the profiler.
@@ -139,7 +143,7 @@ def _profile_ast_tree(self,
         ast.fix_missing_locations(tree)
         return tree
 
-    def profile(self):
+    def profile(self) -> ast.Module:
         """Create an abstract syntax tree of a script and add profiling to it.
 
         Reads a script file and generates an abstract syntax tree.

line_profiler/autoprofile/autoprofile.py

@@ -44,12 +44,15 @@ def main():
     python -m kernprof -p demo.py -l demo.py
     python -m line_profiler -rmt demo.py.lprof
 """
+from __future__ import annotations
 import contextlib
 import functools
 import importlib.util
 import operator
 import sys
 import types
+from collections.abc import MutableMapping
+from typing import Any
 from .ast_tree_profiler import AstTreeProfiler
 from .run_module import AstTreeModuleProfiler
 from .line_profiler_utils import add_imported_function_or_module
@@ -58,7 +61,7 @@ def main():
 PROFILER_LOCALS_NAME = 'prof'
 
 
-def _extend_line_profiler_for_profiling_imports(prof):
+def _extend_line_profiler_for_profiling_imports(prof: Any) -> None:
     """Allow profiler to handle functions/methods, classes & modules with a single call.
 
     Add a method to LineProfiler that can identify whether the object is a
@@ -73,7 +76,9 @@ def _extend_line_profiler_for_profiling_imports(prof):
     prof.add_imported_function_or_module = types.MethodType(add_imported_function_or_module, prof)
 
 
-def run(script_file, ns, prof_mod, profile_imports=False, as_module=False):
+def run(script_file: str, ns: MutableMapping[str, Any],
+        prof_mod: list[str], profile_imports: bool = False,
+        as_module: bool = False) -> None:
     """Automatically profile a script and run it.
 
     Profile functions, classes & modules specified in prof_mod without needing to add
@@ -98,10 +103,10 @@ def run(script_file, ns, prof_mod, profile_imports=False, as_module=False):
             Whether we're running script_file as a module
     """
     class restore_dict:
-        def __init__(self, d, target=None):
+        def __init__(self, d: MutableMapping[str, Any], target=None):
             self.d = d
             self.target = target
-            self.copy = None
+            self.copy: dict[str, Any] | None = None
 
         def __enter__(self):
             assert self.copy is None