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

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,
 )
@@ -647,7 +648,12 @@ def get_meta_prop(o: Union[IMeta, Var]) -> Any:
 _tag_meta = _meta_getter(SYM_TAG_META_KEY)
 
 
-def _loc(form: Union[LispForm, ISeq]) -> Optional[Tuple[int, int, int, int]]:
+T_form = TypeVar("T_form", bound=ReaderForm)
+T_node = TypeVar("T_node", bound=Node)
+LispAnalyzer = Callable[[T_form, AnalyzerContext], T_node]
+
+
+def _loc(form: T_form) -> Optional[Tuple[int, int, int, int]]:
     """Fetch the location of the form in the original filename from the
     input form, if it has metadata."""
     # Technically, IMeta is sufficient for fetching `form.meta` but the
@@ -669,17 +675,17 @@ def _loc(form: Union[LispForm, ISeq]) -> Optional[Tuple[int, int, int, int]]:
     return None
 
 
-def _with_loc(f):
+def _with_loc(f: LispAnalyzer[T_form, T_node]) -> LispAnalyzer[T_form, T_node]:
     """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) -> T_node:
         form_loc = _loc(form)
         if form_loc is None:
             return f(form, ctx)
         else:
-            return f(form, ctx).fix_missing_locations(form_loc)
+            return cast(T_node, f(form, ctx).fix_missing_locations(form_loc))
 
     return _analyze_form
 
@@ -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[T_form, T_node]) -> LispAnalyzer[T_form, T_node]:
     """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) -> T_node:
         assert not ctx.is_quoted, "with-meta nodes are not used in quoted expressions"
 
         descriptor = gen_node(form, ctx)
@@ -825,11 +822,14 @@ def with_meta(
                 assert isinstance(meta_ast, MapNode) or (
                     isinstance(meta_ast, Const) and meta_ast.type == ConstType.MAP
                 )
-                return WithMeta(
-                    form=form,
-                    meta=meta_ast,
-                    expr=descriptor,
-                    env=ctx.get_node_env(pos=ctx.syntax_position),
+                return cast(
+                    T_node,
+                    WithMeta(
+                        form=cast(LispForm, form),
+                        meta=meta_ast,
+                        expr=descriptor,
+                        env=ctx.get_node_env(pos=ctx.syntax_position),
+                    ),
                 )
 
         return descriptor
@@ -3113,7 +3113,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,