Fix various type erros, small refactors

parser/astgen/ast_node.py

@@ -32,8 +32,6 @@ class WalkerCallType(Enum):
 @dataclass
 class AstNode(HasRegion, Generic[MetadataT]):
     region: StrRegion
-    name = None  # type: str
-    del name  # So we get better error msg if we forget to add it to a class
     meta: MetadataT = field(kw_only=True, default=None)
 
     def walk(self, fn: WalkerFnT):

parser/astgen/ast_nodes.py

@@ -10,13 +10,12 @@
     "AstRepeat", "AstIf", "AstWhile", "AstAssign", "AstAugAssign", "AstDefine",
     "AstNumber", "AstString", "AstAnyName", "AstIdent", "AstAttrName",
     "AstListLiteral", "AstAttribute", "AstItem", "AstCall", "AstOp", "AstBinOp",
-    "AstUnaryOp",
+    "AstUnaryOp", "AstDefineParam",
 ]
 
 
 @dataclass
 class AstProgramNode(AstNode[MetadataT]):
-    name = 'program'
     statements: list[AstNode[MetadataT]]
 
     def _walk_members(self, fn: WalkerFnT):
@@ -36,7 +35,6 @@ class VarDeclType(Enum):
 
 @dataclass
 class AstDeclNode(AstNode[MetadataT]):
-    name = 'var_decl'
     scope: VarDeclScope
     type: VarDeclType
     ident: AstIdent[MetadataT]
@@ -48,7 +46,6 @@ def _walk_members(self, fn: WalkerFnT):
 
 @dataclass
 class AstRepeat(AstNode[MetadataT]):
-    name = 'repeat'
     count: AstNode[MetadataT]
     body: list[AstNode[MetadataT]]
 
@@ -58,7 +55,6 @@ def _walk_members(self, fn: WalkerFnT):
 
 @dataclass
 class AstIf(AstNode[MetadataT]):
-    name = 'if'
     cond: AstNode[MetadataT]
     if_body: list[AstNode[MetadataT]]
     # elseif = else{if
@@ -72,7 +68,6 @@ def _walk_members(self, fn: WalkerFnT):
 
 @dataclass
 class AstWhile(AstNode[MetadataT]):
-    name = 'while'
     cond: AstNode[MetadataT]
     body: list[AstNode[MetadataT]]
 
@@ -82,7 +77,6 @@ def _walk_members(self, fn: WalkerFnT):
 
 @dataclass
 class AstAssign(AstNode[MetadataT]):
-    name = '='
     target: AstNode[MetadataT]
     source: AstNode[MetadataT]
 
@@ -96,25 +90,24 @@ class AstAugAssign(AstNode[MetadataT]):
     target: AstNode[MetadataT]
     source: AstNode[MetadataT]
 
-    @property
-    def name(self):
-        return self.op
-
     def _walk_members(self, fn: WalkerFnT):
         self.walk_multiple_objects(fn, (self.target, self.source))
 
 
 @dataclass
 class AstDefine(AstNode[MetadataT]):
-    name = 'def'
-
     ident: AstIdent[MetadataT]
-    # TODO: this should be list[AstDefineParam] where AstParam is an AstNode
-    params: list[tuple[AstIdent[MetadataT], AstIdent[MetadataT]]]  # type, ident
+    params: list[AstDefineParam[MetadataT]]
     body: list[AstNode[MetadataT]]
 
     def _walk_members(self, fn: WalkerFnT):
         self.walk_multiple_objects(fn, (self.ident, self.params, self.body))
+
+
+@dataclass
+class AstDefineParam(AstNode[MetadataT]):
+    type: AstIdent[MetadataT]
+    ident: AstIdent[MetadataT]
 # endregion ---- </Statements> ----
 
 
@@ -135,23 +128,22 @@ class AstAnyName(AstNode[MetadataT]):
     id: str
 
     def __post_init__(self):
-        if type(self) == AstAnyName:
+        if type(self) is AstAnyName:
             raise TypeError("AstAnyName must not be instantiated directly.")
 
 
 @dataclass
 class AstIdent(AstAnyName[MetadataT]):
-    name = 'ident'
+    pass
 
 
 @dataclass
 class AstAttrName(AstAnyName[MetadataT]):
-    name = 'attr'
+    pass
 
 
 @dataclass
 class AstListLiteral(AstNode[MetadataT]):
-    name = 'list'
     items: list[AstNode[MetadataT]]
 
     def _walk_members(self, fn: WalkerFnT):
@@ -160,7 +152,6 @@ def _walk_members(self, fn: WalkerFnT):
 
 @dataclass
 class AstAttribute(AstNode[MetadataT]):
-    name = '.'
     obj: AstNode[MetadataT]
     attr: AstAttrName[MetadataT]
 
@@ -170,7 +161,6 @@ def _walk_members(self, fn: WalkerFnT):
 
 @dataclass
 class AstItem(AstNode[MetadataT]):
-    name = 'item'
     obj: AstNode[MetadataT]
     index: AstNode[MetadataT]
 
@@ -180,7 +170,6 @@ def _walk_members(self, fn: WalkerFnT):
 
 @dataclass
 class AstCall(AstNode[MetadataT]):
-    name = 'call'
     obj: AstNode[MetadataT]
     args: list[AstNode[MetadataT]]
 
@@ -205,10 +194,6 @@ class AstBinOp(AstOp[MetadataT]):
     def __post_init__(self):
         assert self.op in self.valid_ops
 
-    @property
-    def name(self):
-        return self.op
-
     def _walk_members(self, fn: WalkerFnT):
         self.walk_multiple_objects(fn, (self.left, self.right))
 
@@ -222,10 +207,6 @@ class AstUnaryOp(AstOp[MetadataT]):
     def __post_init__(self):
         assert self.op in self.valid_ops
 
-    @property
-    def name(self):
-        return self.op
-
     def _walk_members(self, fn: WalkerFnT):
         self.walk_multiple_objects(fn, (self.operand,))
 # endregion ---- </Expressions> ----

parser/astgen/astgen.py

@@ -4,7 +4,7 @@
 import sys
 from typing import Callable, overload, TypeVar, TypeAlias
 
-from util import flatten_force, is_strict_subclass
+from util import flatten_force, is_strict_subclass, get_mro
 from .ast_nodes import *
 from .eval_literal import eval_number, eval_string
 from .errors import LocatedAstError
@@ -39,11 +39,12 @@
 
 @overload
 def _register_autowalk_expr(node_type: type[AnyNode], /) -> Callable[[CT], CT]: ...
-@overload
+# noinspection overloads
+@overload  # (Pycharm's compatibility detection is not working properly)
 def _register_autowalk_expr(cls: CT, /) -> CT: ...
 
 
-def _register_autowalk_expr(node_type: type[AnyNode] = None, /):
+def _register_autowalk_expr(node_type: type[AnyNode] | None = None, /):
     def decor(fn):
         _AUTOWALK_EXPR_DICT[
             node_type if node_type is not None
@@ -120,11 +121,11 @@ def _walk_smt(self, smt: AnyNode) -> list[AstNode]:
                 smt.region, smt.name, self._walk_assign_left(smt.target),
                 self._walk_expr(smt.source))]
         elif isinstance(smt, DefineNode):
-            # Check arg types during name resolution
-            decls = [(self._walk_ident(d.type), self._walk_ident(d.ident))
-                     for d in smt.args_decl.decls]
-            return [AstDefine(smt.region, self._walk_ident(smt.ident),
-                              decls, self._walk_block(smt.block))]
+            # (Check arg types later during name resolution)
+            return [AstDefine(
+                smt.region, self._walk_ident(smt.ident),
+                [self._walk_arg_decl(a) for a in smt.args_decl.decls],
+                self._walk_block(smt.block))]
         elif isinstance(smt, CallNode):
             # Check that it transforms into smt-intrinsic and
             # not expr-intrinsic in codegen/typecheck
@@ -135,6 +136,9 @@ def _walk_smt(self, smt: AnyNode) -> list[AstNode]:
                 f"expressions have no side-effect so are not allowed at "
                 f"the root level.", smt.region)
 
+    def _walk_arg_decl(self, a: ArgDeclNode):
+        return AstDefineParam(a.region, self._walk_ident(a.type), self._walk_ident(a.ident))
+
     def _walk_var_decl(self, smt: DeclNode):
         scope = (VarDeclScope.LET if isinstance(smt.decl_scope, DeclScope_Let)
                  else VarDeclScope.GLOBAL)
@@ -259,10 +263,10 @@ def autowalk_expr(self, node: AnyNode):
 
     @classmethod
     def _lookup_autowalk_fn(cls, t: type[AnyNode]):
-        assert t != Leaf and t != Node and issubclass(t, AnyNode)
+        assert t is not Leaf and t is not Node and issubclass(t, AnyNode)
         if value := _AUTOWALK_EXPR_DICT.get(t):
             return value
-        for supertype in t.mro():  # See if any supertypes have the autowalker
+        for supertype in get_mro(t):  # See if any supertypes have the autowalker
             if not issubclass(supertype, AnyNode):
                 continue  # Skip - could be a mixin
             if value := _AUTOWALK_EXPR_DICT.get(supertype):

parser/astgen/filtered_walker.py

@@ -4,7 +4,7 @@
 from collections.abc import Callable
 from typing import TypeVar, Generic
 
-from util import flatten_force
+from util import flatten_force, get_mro
 from .ast_node import WalkerCallType, WalkableT, walk_ast
 
 try:
@@ -48,8 +48,8 @@ def __init__(self,
     def copy(self):
         return type(self)(self.enter_cbs, self.exit_cbs, self.both_cbs)
 
-    def instantiate(self, *args: Unpack[Ps]):
-        return BasicFilteredWalker[WT](
+    def instantiate(self, *args: Unpack[Ps]) -> BasicFilteredWalker[WT]:
+        return BasicFilteredWalker(
             self._instantiate_dict(self.enter_cbs, args),
             self._instantiate_dict(self.exit_cbs, args),
             self._instantiate_dict(self.both_cbs, args),
@@ -124,7 +124,7 @@ def __call__(self, o: WalkableT, call_type: WalkerCallType):
     @classmethod
     def _get_funcs(cls, mapping: dict[type[WT] | type, list[VT]], tp: type[WT]) -> list[VT]:
         """Also looks at superclasses/MRO"""
-        return flatten_force([mapping.get(sub, []) for sub in _get_mro(tp)])
+        return flatten_force([mapping.get(sub, []) for sub in get_mro(tp)])
 
 
 class FilteredWalker(BasicFilteredWalker[WT], Generic[WT]):
@@ -173,7 +173,3 @@ def enter_define(self, ...):
         if (parent := fn(cls)) is None:
             return WalkerFilterRegistry()
         return WalkerFilterRegistry.copy(parent)
-
-
-def _get_mro(tp: type) -> tuple[type, ...]:  # tp.__mro__ but with proper types
-    return tp.__mro__  # .mro() recalculates it every time, hence is slow

parser/common/tree_print.py

@@ -183,7 +183,7 @@ def _write_complex_ast_node(self, obj: AstNode, args: list[object],
 
     def _fmt_node_header(self, obj: AnyNode, has_more_args=True):
         args: list[str] = []
-        if type(obj) == Leaf or type(obj) == Node:
+        if type(obj) is Leaf or type(obj) is Node:
             args.append(repr(obj.name))
         if self.verbose:
             args.append(self._fmt_region(obj.region))

parser/cst/named_node.py

@@ -107,7 +107,7 @@ def add(self, *nodes: AnyNode, update_end=False):
 
 
 @overload
-def register_corresponding_token(cls: type[AnyNamedNode],
+def register_corresponding_token(cls: type[AnyNamedNode], /, *,
                                  arity: int | Literal['auto'] | None = None): ...
 
 

parser/cst/token_matcher.py

@@ -150,13 +150,17 @@ def next_idx(self):
     def _match(self):
         if (isinstance(self.pattern, type)
                 and issubclass(self.pattern, Token)
-                and self.pattern != Token):
+                and self.pattern is not Token):
             # optimize the common case of `matches(..., <TokenCls>)`
+            # (Pycharm cannot narrow properly with `and`/`or`)
+            # noinspection bad-argument-type
             self.pattern = _FastTokenClsMatcher(self.pattern)
         elif isinstance(self.pattern, Token) or _issubclass(self.pattern, Token):
+            # (Pycharm once again cannot narrow properly with `and`/`or`)
+            # noinspection bad-argument-type
             self.pattern = TokenMatcher(self.pattern)
         elif isiterable(self.pattern):
-            self.pattern = SeqMatcher(*cast(Iterable[BaseMatcher], self.pattern))
+            self.pattern = SeqMatcher(*self.pattern)
         assert isinstance(self.pattern, BaseMatcher)
         self.result = self.pattern.matches(self.tokens, self.start, self.src)
 

parser/lexer/src_handler.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Sequence, cast
+from typing import Sequence
 
 from util import pack_if_single_item
 from .errors import LocatedTokenizerError

parser/lexer/token_print.py

@@ -7,13 +7,14 @@
 from .tokens import Token
 
 
-def print_tokens(src: str, tokens: list[Token], stream: IO[str] = None, do_ws=False):
+def print_tokens(src: str, tokens: list[Token], stream: IO[str] | None = None,
+                 do_whitespace: bool = False):
     if stream is None:
         stream = sys.stdout
     table = []
     for tok in tokens:
         if tok.is_whitespace:
-            if do_ws:
+            if do_whitespace:
                 table.append(['(WS) ' + repr(tok.region.resolve(src)), tok.name])
         else:
             table.append([str(tok.region.resolve(src)), tok.name])
@@ -23,7 +24,7 @@ def print_tokens(src: str, tokens: list[Token], stream: IO[str] = None, do_ws=Fa
         print(f'{s0:>{max0}} | {s1:>{max1}}', file=stream)
 
 
-def format_tokens(src: str, tokens: list[Token], do_ws=False):
+def format_tokens(src: str, tokens: list[Token], do_whitespace: bool = False):
     out = StringIO()
-    print_tokens(src, tokens, out, do_ws)
+    print_tokens(src, tokens, out, do_whitespace)
     return out.getvalue()

parser/lexer/tokenizer.py

@@ -40,7 +40,7 @@ def prev_content_token_type(self):
         return type(self.content_tokens[-1])
 
     def tokenize(self):
-        idx = 0
+        idx: int = 0
         last_idx = -1
         while idx < len(self.src):
             if idx == last_idx:
@@ -103,9 +103,7 @@ def tokenize(self):
         self.is_done = True
         return self
 
-    def add_token(self, *tokens: Token, whitespace=None) -> int | None:
-        if not tokens:
-            return None
+    def add_token(self, *tokens: Token, whitespace: bool | None = None) -> int:
         self.tokens += tokens
         if whitespace is None:
             self.content_tokens += (t for t in tokens if not t.is_whitespace)