Improve the state of Python type hints in basilisp.lang.*

Author: chrisrink10

CHANGELOG.md

@@ -33,7 +33,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
  * Removed support for PyPy 3.8 (#785)
 
 ### Other
- * Improve the state of the Python type hints in `basilisp.lang.*` (#797)
+ * Improve the state of the Python type hints in `basilisp.lang.*` (#797, #784)
 
 
 ## [v0.1.0b0]

pyproject.toml

@@ -37,6 +37,7 @@ prompt-toolkit = "^3.0.0"
 pyrsistent = "^0.18.0"
 python-dateutil = "^2.8.1"
 readerwriterlock = "^1.0.8"
+typing_extensions = "^4.9.0"
 
 astor = { version = "^0.8.1", python = "<3.9", optional = true }
 pytest = { version = "^7.0.0", optional = true }
@@ -217,6 +218,7 @@ disable = [
 
 [tool.mypy]
 check_untyped_defs = true
+disallow_untyped_decorators = true
 mypy_path = "src/"
 show_error_codes = true
 warn_redundant_casts = true

src/basilisp/lang/atom.py

@@ -1,12 +1,14 @@
 from typing import Callable, Generic, Optional, TypeVar
 
 from readerwriterlock.rwlock import RWLockFair
+from typing_extensions import Concatenate, ParamSpec
 
 from basilisp.lang.interfaces import IPersistentMap, RefValidator
 from basilisp.lang.map import PersistentMap
 from basilisp.lang.reference import RefBase
 
 T = TypeVar("T")
+P = ParamSpec("P")
 
 
 class Atom(RefBase[T], Generic[T]):
@@ -58,7 +60,9 @@ def reset(self, v: T) -> T:
                 self._notify_watches(oldval, v)
                 return v
 
-    def swap(self, f: Callable[..., T], *args, **kwargs) -> T:
+    def swap(
+        self, f: Callable[Concatenate[T, P], T], *args: P.args, **kwargs: P.kwargs
+    ) -> T:
         """Atomically swap the state of the Atom to the return value of
         `f(old, *args, **kwargs)`, returning the new value."""
         while True:

src/basilisp/lang/compiler/analyzer.py

@@ -28,6 +28,7 @@
     Pattern,
     Set,
     Tuple,
+    TypeVar,
     Union,
     cast,
 )
@@ -669,12 +670,17 @@ def _loc(form: Union[LispForm, ISeq]) -> Optional[Tuple[int, int, int, int]]:
     return None
 
 
-def _with_loc(f):
+T_form = TypeVar("T_form", bound=LispForm)
+T_node = TypeVar("T_node", Node, None)
+LispAnalyzer = Callable[[T_form, AnalyzerContext], T_node]
+
+
+def _with_loc(f: LispAnalyzer) -> LispAnalyzer:
     """Attach any available location information from the input form to
     the node environment returned from the parsing function."""
 
     @wraps(f)
-    def _analyze_form(form: Union[LispForm, ISeq], ctx: AnalyzerContext) -> Node:
+    def _analyze_form(form: T_form, ctx: AnalyzerContext) -> Node:
         form_loc = _loc(form)
         if form_loc is None:
             return f(form, ctx)
@@ -795,24 +801,15 @@ def _tag_ast(form: Optional[LispForm], ctx: AnalyzerContext) -> Optional[Node]:
     return _analyze_form(form, ctx)
 
 
-def _with_meta(gen_node):
+def _with_meta(gen_node: LispAnalyzer) -> LispAnalyzer:
     """Wraps the node generated by gen_node in a :with-meta AST node if the
     original form has meta.
 
     :with-meta AST nodes are used for non-quoted collection literals and for
     function expressions."""
 
     @wraps(gen_node)
-    def with_meta(
-        form: Union[
-            llist.PersistentList,
-            lmap.PersistentMap,
-            ISeq,
-            lset.PersistentSet,
-            vec.PersistentVector,
-        ],
-        ctx: AnalyzerContext,
-    ) -> Node:
+    def with_meta(form: T_form, ctx: AnalyzerContext) -> Node:
         assert not ctx.is_quoted, "with-meta nodes are not used in quoted expressions"
 
         descriptor = gen_node(form, ctx)
@@ -826,7 +823,7 @@ def with_meta(
                     isinstance(meta_ast, Const) and meta_ast.type == ConstType.MAP
                 )
                 return WithMeta(
-                    form=form,
+                    form=cast(T_form, form),
                     meta=meta_ast,
                     expr=descriptor,
                     env=ctx.get_node_env(pos=ctx.syntax_position),
@@ -3113,7 +3110,7 @@ def _yield_ast(form: ISeq, ctx: AnalyzerContext) -> Yield:
         return Yield.expressionless(form, ctx.get_node_env(pos=ctx.syntax_position))
 
 
-SpecialFormHandler = Callable[[ISeq, AnalyzerContext], SpecialFormNode]
+SpecialFormHandler = Callable[[T_form, AnalyzerContext], SpecialFormNode]
 _SPECIAL_FORM_HANDLERS: Mapping[sym.Symbol, SpecialFormHandler] = {
     SpecialForm.AWAIT: _await_ast,
     SpecialForm.DEF: _def_ast,

src/basilisp/lang/compiler/generator.py

@@ -31,6 +31,7 @@
 )
 
 import attr
+from typing_extensions import Concatenate, ParamSpec
 
 from basilisp.lang import keyword as kw
 from basilisp.lang import list as llist
@@ -93,9 +94,9 @@
     PyTuple,
 )
 from basilisp.lang.compiler.nodes import Queue as QueueNode
-from basilisp.lang.compiler.nodes import Quote, ReaderLispForm, Recur, Reify, Require
+from basilisp.lang.compiler.nodes import Quote, Recur, Reify, Require
 from basilisp.lang.compiler.nodes import Set as SetNode
-from basilisp.lang.compiler.nodes import SetBang, Throw, Try, VarRef
+from basilisp.lang.compiler.nodes import SetBang, T_withmeta, Throw, Try, VarRef
 from basilisp.lang.compiler.nodes import Vector as VectorNode
 from basilisp.lang.compiler.nodes import WithMeta, Yield
 from basilisp.lang.interfaces import IMeta, IRecord, ISeq, ISeqable, IType
@@ -332,8 +333,11 @@ def reduce(*genned: "GeneratedPyAST[T_pynode]") -> "GeneratedPyAST[T_pynode]":
 
 
 PyASTStream = Iterable[ast.AST]
-SimplePyASTGenerator = Callable[[GeneratorContext, ReaderLispForm], GeneratedPyAST]
-PyASTGenerator = Callable[[GeneratorContext, Node], GeneratedPyAST]
+T_node = TypeVar("T_node", bound=Node)
+P_generator = ParamSpec("P_generator")
+PyASTGenerator = Callable[
+    Concatenate[GeneratorContext, T_node, P_generator], GeneratedPyAST[T_pynode]
+]
 
 
 ####################
@@ -369,12 +373,19 @@ def attr_node(node, idx):
     return attr_node(ast.Name(id=attrs[0], ctx=ast.Load()), 1)
 
 
-def _simple_ast_generator(gen_ast):
+P_simplegen = ParamSpec("P_simplegen")
+
+
+def _simple_ast_generator(
+    gen_ast: Callable[P_simplegen, ast.AST]
+) -> Callable[P_simplegen, GeneratedPyAST]:
     """Wrap simpler AST generators to return a GeneratedPyAST."""
 
     @wraps(gen_ast)
-    def wrapped_ast_generator(ctx: GeneratorContext, form: LispForm) -> GeneratedPyAST:
-        return GeneratedPyAST(node=gen_ast(ctx, form))
+    def wrapped_ast_generator(
+        *args: P_simplegen.args, **kwargs: P_simplegen.kwargs
+    ) -> GeneratedPyAST:
+        return GeneratedPyAST(node=gen_ast(*args, **kwargs))
 
     return wrapped_ast_generator
 
@@ -524,23 +535,30 @@ def _ast_with_loc(
     return py_ast
 
 
-def _with_ast_loc(f):
+def _with_ast_loc(
+    f: PyASTGenerator[T_node, P_generator, T_pynode]
+) -> PyASTGenerator[T_node, P_generator, T_pynode]:
     """Wrap a generator function in a decorator to supply line and column
     information to the returned Python AST node. Dependency nodes will not
     be hydrated, functions whose returns need dependency nodes to be
     hydrated should use `_with_ast_loc_deps` below."""
 
     @wraps(f)
     def with_lineno_and_col(
-        ctx: GeneratorContext, node: Node, *args, **kwargs
-    ) -> GeneratedPyAST:
+        ctx: GeneratorContext,
+        node: T_node,
+        *args: P_generator.args,
+        **kwargs: P_generator.kwargs,
+    ) -> GeneratedPyAST[T_pynode]:
         py_ast = f(ctx, node, *args, **kwargs)
         return _ast_with_loc(py_ast, node.env)
 
     return with_lineno_and_col
 
 
-def _with_ast_loc_deps(f):
+def _with_ast_loc_deps(
+    f: PyASTGenerator[T_node, P_generator, T_pynode]
+) -> PyASTGenerator[T_node, P_generator, T_pynode]:
     """Wrap a generator function in a decorator to supply line and column
     information to the returned Python AST node and dependency nodes.
 
@@ -551,8 +569,11 @@ def _with_ast_loc_deps(f):
 
     @wraps(f)
     def with_lineno_and_col(
-        ctx: GeneratorContext, node: Node, *args, **kwargs
-    ) -> GeneratedPyAST:
+        ctx: GeneratorContext,
+        node: T_node,
+        *args: P_generator.args,
+        **kwargs: P_generator.kwargs,
+    ) -> GeneratedPyAST[T_pynode]:
         py_ast = f(ctx, node, *args, **kwargs)
         return _ast_with_loc(py_ast, node.env, include_dependencies=True)
 
@@ -847,16 +868,15 @@ def _def_to_py_ast(  # pylint: disable=too-many-locals
         assert node.init is not None  # silence MyPy
         if node.init.op == NodeOp.FN:
             assert isinstance(node.init, Fn)
-            def_ast = _fn_to_py_ast(  # type: ignore[call-arg]
-                ctx, node.init, def_name=defsym.name
-            )
+            def_ast = _fn_to_py_ast(ctx, node.init, def_name=defsym.name)
             is_defn = True
         elif (
             node.init.op == NodeOp.WITH_META
             and isinstance(node.init, WithMeta)
             and node.init.expr.op == NodeOp.FN
         ):
             assert isinstance(node.init, WithMeta)
+            assert isinstance(node.init.expr, Fn)
             def_ast = _with_meta_to_py_ast(ctx, node.init, def_name=defsym.name)
             is_defn = True
         else:
@@ -1315,10 +1335,8 @@ def __deftype_staticmethod_to_py_ast(
         )
 
 
-DefTypeASTGenerator = Callable[
-    [GeneratorContext, DefTypeMember], GeneratedPyAST[ast.stmt]
-]
-_DEFTYPE_MEMBER_HANDLER: Mapping[NodeOp, DefTypeASTGenerator] = {
+T_deftypenode = TypeVar("T_deftypenode", bound=DefTypeMember)
+_DEFTYPE_MEMBER_HANDLER: Mapping[NodeOp, PyASTGenerator] = {
     NodeOp.DEFTYPE_CLASSMETHOD: __deftype_classmethod_to_py_ast,
     NodeOp.DEFTYPE_METHOD: __deftype_method_to_py_ast,
     NodeOp.DEFTYPE_PROPERTY: __deftype_property_to_py_ast,
@@ -1328,7 +1346,7 @@ def __deftype_staticmethod_to_py_ast(
 
 def __deftype_member_to_py_ast(
     ctx: GeneratorContext,
-    node: DefTypeMember,
+    node: T_deftypenode,
 ) -> GeneratedPyAST[ast.stmt]:
     member_type = node.op
     handle_deftype_member = _DEFTYPE_MEMBER_HANDLER.get(member_type)
@@ -1485,7 +1503,9 @@ def _deftype_to_py_ast(  # pylint: disable=too-many-locals
         )
 
 
-def _wrap_override_var_indirection(f: PyASTGenerator) -> PyASTGenerator:
+def _wrap_override_var_indirection(
+    f: PyASTGenerator[T_node, P_generator, T_pynode]
+) -> PyASTGenerator[T_node, P_generator, T_pynode]:
     """
     Wrap a Node generator to apply a special override requiring Var indirection
     for any Var accesses generated within top-level `do` blocks.
@@ -1504,11 +1524,14 @@ def _wrap_override_var_indirection(f: PyASTGenerator) -> PyASTGenerator:
 
     @wraps(f)
     def _wrapped_do(
-        ctx: GeneratorContext, node: Node, *args, **kwargs
-    ) -> GeneratedPyAST:
+        ctx: GeneratorContext,
+        node: T_node,
+        *args: P_generator.args,
+        **kwargs: P_generator.kwargs,
+    ) -> GeneratedPyAST[T_pynode]:
         if isinstance(node, Do) and node.top_level:
             with ctx.with_var_indirection_override():
-                return f(ctx, node, *args, **kwargs)
+                return f(ctx, cast(T_node, node), *args, **kwargs)
         else:
             with ctx.with_var_indirection_override(False):
                 return f(ctx, node, *args, **kwargs)
@@ -3358,7 +3381,7 @@ def _py_tuple_to_py_ast(
 ############
 
 
-_WITH_META_EXPR_HANDLER = {
+_WITH_META_EXPR_HANDLER: Mapping[NodeOp, PyASTGenerator] = {
     NodeOp.FN: _fn_to_py_ast,
     NodeOp.MAP: _map_to_py_ast,
     NodeOp.QUEUE: _queue_to_py_ast,
@@ -3369,7 +3392,10 @@ def _py_tuple_to_py_ast(
 
 
 def _with_meta_to_py_ast(
-    ctx: GeneratorContext, node: WithMeta, **kwargs
+    ctx: GeneratorContext,
+    node: WithMeta[T_withmeta],
+    *args: P_generator.args,
+    **kwargs: P_generator.kwargs,
 ) -> GeneratedPyAST[ast.expr]:
     """Generate a Python AST node for Python interop method calls."""
     assert node.op == NodeOp.WITH_META
@@ -3378,7 +3404,7 @@ def _with_meta_to_py_ast(
     assert (
         handle_expr is not None
     ), "No expression handler for with-meta child node type"
-    return handle_expr(ctx, node.expr, meta_node=node.meta, **kwargs)
+    return handle_expr(ctx, node.expr, meta_node=node.meta, *args, **kwargs)
 
 
 #################
@@ -3634,7 +3660,7 @@ def _const_record_to_py_ast(
         key_nodes = _kw_to_py_ast(k, ctx)
         keys.append(key_nodes.node)
         assert (
-            len(key_nodes.dependencies) == 0
+            not key_nodes.dependencies
         ), "Simple AST generators must emit no dependencies"
 
         val_nodes = _const_val_to_py_ast(v, ctx)
@@ -3759,7 +3785,7 @@ def _const_node_to_py_ast(
     NodeOp.PY_TUPLE: _py_tuple_to_py_ast,
     NodeOp.QUEUE: _queue_to_py_ast,
     NodeOp.QUOTE: _quote_to_py_ast,
-    NodeOp.RECUR: _recur_to_py_ast,  # type: ignore
+    NodeOp.RECUR: _recur_to_py_ast,
     NodeOp.REIFY: _reify_to_py_ast,
     NodeOp.REQUIRE: _require_to_py_ast,
     NodeOp.SET: _set_to_py_ast,
@@ -3769,7 +3795,7 @@ def _const_node_to_py_ast(
     NodeOp.YIELD: _yield_to_py_ast,
     NodeOp.VAR: _var_sym_to_py_ast,
     NodeOp.VECTOR: _vec_to_py_ast,
-    NodeOp.WITH_META: _with_meta_to_py_ast,  # type: ignore
+    NodeOp.WITH_META: _with_meta_to_py_ast,
 }
 
 

src/basilisp/lang/futures.py

@@ -5,10 +5,12 @@
 from typing import Callable, Optional, TypeVar
 
 import attr
+from typing_extensions import ParamSpec
 
 from basilisp.lang.interfaces import IBlockingDeref
 
 T = TypeVar("T")
+P = ParamSpec("P")
 
 
 @attr.frozen(eq=True, repr=False)
@@ -55,8 +57,8 @@ def __init__(self, max_workers: Optional[int] = None):
         super().__init__(max_workers=max_workers)
 
     # pylint: disable=arguments-differ
-    def submit(  # type: ignore
-        self, fn: Callable[..., T], *args, **kwargs
+    def submit(  # type: ignore[override]
+        self, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs
     ) -> "Future[T]":
         return Future(super().submit(fn, *args, **kwargs))
 
@@ -70,7 +72,7 @@ def __init__(
         super().__init__(max_workers=max_workers, thread_name_prefix=thread_name_prefix)
 
     # pylint: disable=arguments-differ
-    def submit(  # type: ignore
-        self, fn: Callable[..., T], *args, **kwargs
+    def submit(  # type: ignore[override]
+        self, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs
     ) -> "Future[T]":
         return Future(super().submit(fn, *args, **kwargs))