jaclang 0.7.0__py3-none-any.whl → 0.7.2__py3-none-any.whl

jaclang/compiler/absyntree.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import ast as ast3
+import os
 from hashlib import md5
 from types import EllipsisType
 from typing import Any, Callable, Generic, Optional, Sequence, Type, TypeVar
@@ -31,7 +32,6 @@ class AstNode:
         self.kid: list[AstNode] = [x.set_parent(self) for x in kid]
         self.sym_tab: Optional[SymbolTable] = None
         self._sub_node_tab: dict[type, list[AstNode]] = {}
-        self._typ: type = type(None)
         self.gen: CodeGenTarget = CodeGenTarget()
         self.meta: dict[str, str] = {}
         self.loc: CodeLocInfo = CodeLocInfo(*self.resolve_tok_range())
@@ -111,11 +111,15 @@ class AstNode:
 
         return Pass.get_all_sub_nodes(node=self, typ=typ, brute_force=brute_force)
 
-    def parent_of_type(self, typ: Type[T]) -> T:
+    def find_parent_of_type(self, typ: Type[T]) -> Optional[T]:
         """Get parent of type."""
         from jaclang.compiler.passes import Pass
 
-        ret = Pass.has_parent_of_type(node=self, typ=typ)
+        return Pass.has_parent_of_type(node=self, typ=typ)
+
+    def parent_of_type(self, typ: Type[T]) -> T:
+        """Get parent of type."""
+        ret = self.find_parent_of_type(typ)
         if isinstance(ret, typ):
             return ret
         else:
@@ -178,6 +182,8 @@ class AstSymbolNode(AstNode):
         self.sym_link: Optional[Symbol] = None
         self.sym_name: str = sym_name
         self.sym_name_node = sym_name_node
+        if isinstance(self.sym_name_node, NameSpec):
+            self.sym_name_node.name_of = self
         self.sym_type: SymbolType = sym_type
         self.sym_info: SymbolInfo = SymbolInfo()
         self.py_ctx_func: Type[ast3.AST] = ast3.Load
@@ -258,8 +264,11 @@ class WalkerStmtOnlyNode(AstNode):
 class AstImplOnlyNode(AstNode):
     """ImplOnly node type for Jac Ast."""
 
-    def __init__(self, body: SubNodeList, decl_link: Optional[AstNode]) -> None:
+    def __init__(
+        self, target: ArchRefChain, body: SubNodeList, decl_link: Optional[AstNode]
+    ) -> None:
         """Initialize impl only node."""
+        self.target = target
         self.body = body
         self.decl_link = decl_link
 
@@ -304,6 +313,10 @@ class CodeBlockStmt(AstNode):
 class NameSpec(AtomExpr, EnumBlockStmt):
     """NameSpec node type for Jac Ast."""
 
+    def __init__(self) -> None:
+        """Initialize name spec node."""
+        self.name_of: AstSymbolNode = self
+
 
 class ArchSpec(ElementStmt, CodeBlockStmt, AstSymbolNode, AstDocNode, AstSemStrNode):
     """ArchSpec node type for Jac Ast."""
@@ -392,7 +405,7 @@ class Module(AstDocNode):
         body: Sequence[ElementStmt | String | EmptyToken],
         is_imported: bool,
         kid: Sequence[AstNode],
-        test_mod: Optional[Module] = None,
+        stub_only: bool = False,
         registry: Optional[SemRegistry] = None,
     ) -> None:
         """Initialize whole program node."""
@@ -400,13 +413,37 @@ class Module(AstDocNode):
         self.source = source
         self.body = body
         self.is_imported = is_imported
+        self.stub_only = stub_only
         self.impl_mod: list[Module] = []
-        self.test_mod = test_mod
+        self.test_mod: list[Module] = []
         self.mod_deps: dict[str, Module] = {}
         self.registry = registry
         AstNode.__init__(self, kid=kid)
         AstDocNode.__init__(self, doc=doc)
 
+    @property
+    def annexable_by(self) -> Optional[str]:
+        """Get annexable by."""
+        if not self.stub_only and self.loc.mod_path.endswith("impl.jac"):
+            head_mod_name = self.name.split(".")[0]
+            potential_path = os.path.join(
+                os.path.dirname(self.loc.mod_path),
+                f"{head_mod_name}.jac",
+            )
+            if os.path.exists(potential_path):
+                return potential_path
+            if os.path.split(os.path.dirname(self.loc.mod_path))[-1].endswith(".impl"):
+                head_mod_name = os.path.split(os.path.dirname(self.loc.mod_path))[
+                    -1
+                ].split(".")[0]
+                potential_path = os.path.join(
+                    os.path.dirname(os.path.dirname(self.loc.mod_path)),
+                    f"{head_mod_name}.jac",
+                )
+                if os.path.exists(potential_path):
+                    return potential_path
+        return None
+
     def normalize(self, deep: bool = False) -> bool:
         """Normalize module node."""
         res = True
@@ -701,10 +738,12 @@ class ModuleItem(AstSymbolNode):
         name: Name,
         alias: Optional[Name],
         kid: Sequence[AstNode],
+        sub_module: Optional[Module] = None,
     ) -> None:
         """Initialize module item node."""
         self.name = name
         self.alias = alias
+        self.sub_module = sub_module
         AstNode.__init__(self, kid=kid)
         AstSymbolNode.__init__(
             self,
@@ -839,7 +878,6 @@ class ArchDef(ArchSpec, AstImplOnlyNode):
         decl_link: Optional[Architype] = None,
     ) -> None:
         """Initialize arch def node."""
-        self.target = target
         AstNode.__init__(self, kid=kid)
         AstSymbolNode.__init__(
             self,
@@ -849,7 +887,7 @@ class ArchDef(ArchSpec, AstImplOnlyNode):
         )
         AstDocNode.__init__(self, doc=doc)
         ArchSpec.__init__(self, decorators=decorators)
-        AstImplOnlyNode.__init__(self, body=body, decl_link=decl_link)
+        AstImplOnlyNode.__init__(self, target=target, body=body, decl_link=decl_link)
 
     def normalize(self, deep: bool = False) -> bool:
         """Normalize arch def node."""
@@ -950,7 +988,6 @@ class EnumDef(ArchSpec, AstImplOnlyNode):
         decl_link: Optional[Enum] = None,
     ) -> None:
         """Initialize arch def node."""
-        self.target = target
         AstNode.__init__(self, kid=kid)
         AstSymbolNode.__init__(
             self,
@@ -960,7 +997,7 @@ class EnumDef(ArchSpec, AstImplOnlyNode):
         )
         AstDocNode.__init__(self, doc=doc)
         ArchSpec.__init__(self, decorators=decorators)
-        AstImplOnlyNode.__init__(self, body=body, decl_link=decl_link)
+        AstImplOnlyNode.__init__(self, target=target, body=body, decl_link=decl_link)
 
     def normalize(self, deep: bool = False) -> bool:
         """Normalize enum def node."""
@@ -1000,7 +1037,7 @@ class Ability(
         is_static: bool,
         is_abstract: bool,
         access: Optional[SubTag[Token]],
-        signature: Optional[FuncSignature | EventSignature],
+        signature: FuncSignature | EventSignature,
         body: Optional[SubNodeList[CodeBlockStmt] | AbilityDef | FuncCall],
         kid: Sequence[AstNode],
         semstr: Optional[String] = None,
@@ -1027,16 +1064,6 @@ class Ability(
         AstDocNode.__init__(self, doc=doc)
         AstAsyncNode.__init__(self, is_async=is_async)
 
-    @property
-    def is_method(self) -> bool:
-        """Check if is method."""
-        check = isinstance(self.parent, SubNodeList) and isinstance(
-            self.parent.parent, Architype
-        )
-        if check:
-            self.sym_type = SymbolType.METHOD
-        return check
-
     @property
     def is_func(self) -> bool:
         """Check if is func."""
@@ -1119,7 +1146,6 @@ class AbilityDef(AstSymbolNode, ElementStmt, AstImplOnlyNode, CodeBlockStmt):
         decl_link: Optional[Ability] = None,
     ) -> None:
         """Initialize ability def node."""
-        self.target = target
         self.signature = signature
         self.decorators = decorators
         AstNode.__init__(self, kid=kid)
@@ -1130,7 +1156,7 @@ class AbilityDef(AstSymbolNode, ElementStmt, AstImplOnlyNode, CodeBlockStmt):
             sym_type=SymbolType.IMPL,
         )
         AstDocNode.__init__(self, doc=doc)
-        AstImplOnlyNode.__init__(self, body=body, decl_link=decl_link)
+        AstImplOnlyNode.__init__(self, target=target, body=body, decl_link=decl_link)
 
     def normalize(self, deep: bool = False) -> bool:
         """Normalize ability def node."""
@@ -1152,14 +1178,6 @@ class AbilityDef(AstSymbolNode, ElementStmt, AstImplOnlyNode, CodeBlockStmt):
         self.set_kids(nodes=new_kid)
         return res
 
-    @property
-    def is_method(self) -> bool:
-        """Check if is method."""
-        return (
-            len(self.target.archs) > 1
-            and self.target.archs[-2].arch.name != Tok.ABILITY_OP
-        )
-
 
 class FuncSignature(AstSemStrNode):
     """FuncSignature node type for Jac Ast."""
@@ -1197,13 +1215,6 @@ class FuncSignature(AstSemStrNode):
         self.set_kids(nodes=new_kid)
         return res
 
-    @property
-    def is_method(self) -> bool:
-        """Check if is method."""
-        return (isinstance(self.parent, Ability) and self.parent.is_method) or (
-            isinstance(self.parent, AbilityDef) and self.parent.is_method
-        )
-
     @property
     def is_static(self) -> bool:
         """Check if is static."""
@@ -1256,17 +1267,6 @@ class EventSignature(AstSemStrNode):
         self.set_kids(nodes=new_kid)
         return res
 
-    @property
-    def is_method(self) -> bool:
-        """Check if is method."""
-        if (isinstance(self.parent, Ability) and self.parent.is_method) or (
-            isinstance(self.parent, AbilityDef)
-            and isinstance(self.parent.decl_link, Ability)
-            and self.parent.decl_link.is_method
-        ):
-            return True
-        return False
-
 
 class ArchRefChain(AstNode):
     """Arch ref list node type for Jac Ast."""
@@ -1428,7 +1428,7 @@ class HasVar(AstSymbolNode, AstTypedVarNode, AstSemStrNode):
             self,
             sym_name=name.value,
             sym_name_node=name,
-            sym_type=SymbolType.VAR,
+            sym_type=SymbolType.HAS_VAR,
         )
         AstTypedVarNode.__init__(self, type_tag=type_tag)
         AstSemStrNode.__init__(self, semstr=semstr)
@@ -3100,6 +3100,7 @@ class ArchRef(NameSpec):
             sym_name_node=name_ref,
             sym_type=SymbolType.TYPE,
         )
+        NameSpec.__init__(self)
 
     def normalize(self, deep: bool = False) -> bool:
         """Normalize ast node."""
@@ -3137,6 +3138,7 @@ class SpecialVarRef(NameSpec):
             sym_name_node=var,
             sym_type=SymbolType.VAR,
         )
+        NameSpec.__init__(self)
 
     def normalize(self, deep: bool = False) -> bool:
         """Normalize ast node."""
@@ -3820,6 +3822,7 @@ class Name(Token, NameSpec):
             sym_name_node=self,
             sym_type=SymbolType.VAR,
         )
+        NameSpec.__init__(self)
 
     def unparse(self) -> str:
         """Unparse name."""

jaclang/compiler/compile.py

@@ -57,6 +57,27 @@ def jac_str_to_pass(
     return ast_ret
 
 
+def jac_ir_to_pass(
+    ir: ast.AstNode,
+    target: Optional[Type[Pass]] = None,
+    schedule: list[Type[Pass]] = pass_schedule,
+) -> Pass:
+    """Convert a Jac file to an AST."""
+    if not target:
+        target = schedule[-1] if schedule else None
+    ast_ret = (
+        Pass(input_ir=ir, prior=None)
+        if not len(schedule)
+        else schedule[0](input_ir=ir, prior=None)
+    )
+    for i in schedule[1:]:
+        if i == target:
+            break
+        ast_ret = i(input_ir=ast_ret.ir, prior=ast_ret)
+    ast_ret = target(input_ir=ast_ret.ir, prior=ast_ret) if target else ast_ret
+    return ast_ret
+
+
 def jac_pass_to_pass(
     in_pass: Pass,
     target: Optional[Type[Pass]] = None,

jaclang/compiler/passes/main/__init__.py

@@ -3,7 +3,7 @@
 from .sub_node_tab_pass import SubNodeTabPass
 from .import_pass import JacImportPass, PyImportPass  # noqa: I100
 from .sym_tab_build_pass import SymTabBuildPass  # noqa: I100
-from .def_impl_match_pass import DeclDefMatchPass  # noqa: I100
+from .def_impl_match_pass import DeclImplMatchPass  # noqa: I100
 from .def_use_pass import DefUsePass  # noqa: I100
 from .pyout_pass import PyOutPass  # noqa: I100
 from .pyast_load_pass import PyastBuildPass  # type: ignore # noqa: I100
@@ -20,7 +20,7 @@ __all__ = [
     "JacImportPass",
     "PyImportPass",
     "SymTabBuildPass",
-    "DeclDefMatchPass",
+    "DeclImplMatchPass",
     "DefUsePass",
     "PyOutPass",
     "PyastBuildPass",

jaclang/compiler/passes/main/def_impl_match_pass.py

@@ -9,10 +9,10 @@ body field.
 import jaclang.compiler.absyntree as ast
 from jaclang.compiler.passes import Pass
 from jaclang.compiler.passes.main import SubNodeTabPass
-from jaclang.compiler.symtable import Symbol, SymbolTable, SymbolType
+from jaclang.compiler.symtable import Symbol, SymbolTable
 
 
-class DeclDefMatchPass(Pass):
+class DeclImplMatchPass(Pass):
     """Decls and Def matching pass."""
 
     def enter_module(self, node: ast.Module) -> None:
@@ -39,11 +39,12 @@ class DeclDefMatchPass(Pass):
     def connect_def_impl(self, sym_tab: SymbolTable) -> None:
         """Connect Decls and Defs."""
         for sym in sym_tab.tab.values():
-            if sym.sym_type == SymbolType.IMPL:
+            if isinstance(sym.decl, ast.AstImplOnlyNode):
                 # currently strips the type info from impls
                 arch_refs = [x[3:] for x in sym.sym_name.split(".")]
+                name_of_links = []  # to link archref names to decls
                 lookup = sym_tab.lookup(arch_refs[0])
-                # Below may need to be a while instead of if to skip over local
+                # If below may need to be a while instead of if to skip over local
                 # import name collisions (see test: test_impl_decl_resolution_fix)
                 if lookup and not isinstance(lookup.decl, ast.AstImplNeedingNode):
                     lookup = sym_tab.parent.lookup(arch_refs[0])
@@ -52,6 +53,7 @@ class DeclDefMatchPass(Pass):
                     if len(arch_refs) == 1 and lookup
                     else lookup.defn[-1] if lookup else None
                 )
+                name_of_links.append(decl_node) if decl_node else None
                 for name in arch_refs[1:]:
                     if decl_node:
                         lookup = (
@@ -64,6 +66,7 @@ class DeclDefMatchPass(Pass):
                             if len(arch_refs) == 1 and lookup
                             else lookup.defn[-1] if lookup else None
                         )
+                        name_of_links.append(decl_node) if decl_node else None
                     else:
                         break
                 if not decl_node:
@@ -76,10 +79,9 @@ class DeclDefMatchPass(Pass):
                     continue
                 decl_node.body = sym.decl  # type: ignore
                 sym.decl.decl_link = decl_node  # type: ignore
-                source_node = sym.decl.parent
-                decl_node.add_kids_right([sym.decl])  # type: ignore
-                if source_node and sym.decl in source_node.kid:
-                    source_node.kid.remove(sym.decl)
+                for idx, a in enumerate(sym.decl.target.archs):
+                    if isinstance(a.name_ref.sym_name_node, ast.NameSpec):
+                        a.name_ref.sym_name_node.name_of = name_of_links[idx]
                 decl_node.sym_tab.tab = sym.decl.sym_tab.tab  # type: ignore
         for i in sym_tab.kid:
             self.connect_def_impl(i)

jaclang/compiler/passes/main/def_use_pass.py

@@ -28,13 +28,20 @@ class DefUsePass(SymTabPass):
         sym_tab: Optional[SymbolTable],
         """
         if node.arch_type.name == Tok.KW_WALKER:
-            for i in (
-                self.get_all_sub_nodes(node, ast.VisitStmt)
-                + self.get_all_sub_nodes(node, ast.IgnoreStmt)
-                + self.get_all_sub_nodes(node, ast.DisengageStmt)
-                + self.get_all_sub_nodes(node, ast.EdgeOpRef)
-            ):
-                i.from_walker = True
+            self.inform_from_walker(node)
+            for i in self.get_all_sub_nodes(node, ast.Ability):
+                if isinstance(i.body, ast.AbilityDef):
+                    self.inform_from_walker(i.body)
+
+    def inform_from_walker(self, node: ast.AstNode) -> None:
+        """Inform all sub nodes that they are from a walker."""
+        for i in (
+            self.get_all_sub_nodes(node, ast.VisitStmt)
+            + self.get_all_sub_nodes(node, ast.IgnoreStmt)
+            + self.get_all_sub_nodes(node, ast.DisengageStmt)
+            + self.get_all_sub_nodes(node, ast.EdgeOpRef)
+        ):
+            i.from_walker = True
 
     def enter_arch_ref(self, node: ast.ArchRef) -> None:
         """Sub objects.

jaclang/compiler/passes/main/fuse_typeinfo_pass.py

@@ -129,8 +129,10 @@ class FuseTypeInfoPass(Pass):
         if isinstance(mypy_node, MypyNodes.MemberExpr):
             if mypy_node in self.node_type_hash:
                 t = str(self.node_type_hash[mypy_node])
-                if "->" in t:
+                if "def" in t and "->" in t:
                     t = t.split("->")[1].strip()
+                elif "def" in t:
+                    t = "None"
                 node.sym_info.typ = t
             else:
                 self.__debug_print(f"{node.loc} MemberExpr type is not found")
@@ -181,6 +183,18 @@ class FuseTypeInfoPass(Pass):
         """Pass handler for name nodes."""
         self.__collect_type_from_symbol(node)
 
+        # Assign correct symbols to sym_link in case of
+        # AtomTrailer Object
+        if isinstance(node.parent, ast.AtomTrailer):
+            target_node = node.parent.target
+            if isinstance(target_node, ast.AstSymbolNode):
+                parent_symbol_table = target_node.sym_info.typ_sym_table
+                if isinstance(parent_symbol_table, ast.SymbolTable):
+                    owner = parent_symbol_table.owner
+                    if isinstance(owner, ast.AstSymbolNode):
+                        target_node.sym_link = owner.sym_link
+                        node.sym_link = parent_symbol_table.lookup(node.sym_name)
+
     @__handle_node
     def enter_module_path(self, node: ast.ModulePath) -> None:
         """Pass handler for ModulePath nodes."""
@@ -263,6 +277,11 @@ class FuseTypeInfoPass(Pass):
                 "Getting type of 'HasVar' is only supported with AssignmentStmt"
             )
 
+    def exit_has_var(self, node: ast.HasVar) -> None:
+        """Pass handler for HasVar nodes."""
+        node.sym_info.typ = node.name.sym_info.typ
+        node.sym_info.typ_sym_table = node.name.sym_info.typ_sym_table
+
     @__handle_node
     def enter_multi_string(self, node: ast.MultiString) -> None:
         """Pass handler for MultiString nodes."""

jaclang/compiler/passes/main/import_pass.py

@@ -46,17 +46,25 @@ class JacImportPass(Pass):
         """Process an import."""
         lang = i.parent_of_type(ast.Import).hint.tag.value
         if lang == "jac" and not i.sub_module:
-            mod = self.import_jac_module(
+            self.import_jac_module(
                 node=i,
                 mod_path=node.loc.mod_path,
             )
-            if mod:
-                self.run_again = True
-                i.sub_module = mod
-                i.add_kids_right([mod], pos_update=False)
+
+    def attach_mod_to_node(
+        self, node: ast.ModulePath | ast.ModuleItem, mod: ast.Module | None
+    ) -> None:
+        """Attach a module to a node."""
+        if mod:
+            self.run_again = True
+            node.sub_module = mod
+            self.__annex_impl(mod)
+            node.add_kids_right([mod], pos_update=False)
 
     def __annex_impl(self, node: ast.Module) -> None:
         """Annex impl and test modules."""
+        if node.stub_only:
+            return
         if not node.loc.mod_path:
             self.error("Module has no path")
         if not node.loc.mod_path.endswith(".jac"):
@@ -99,7 +107,7 @@ class JacImportPass(Pass):
             ) and cur_file.endswith(".test.jac"):
                 mod = self.import_jac_mod_from_file(cur_file)
                 if mod:
-                    node.test_mod = mod
+                    node.test_mod.append(mod)
                     node.add_kids_right([mod], pos_update=False)
                     mod.parent = node
 
@@ -114,15 +122,56 @@ class JacImportPass(Pass):
             node.sub_module.name = node.alias.value
         # Items matched during def/decl pass
 
-    def import_jac_module(
-        self, node: ast.ModulePath, mod_path: str
-    ) -> ast.Module | None:
+    def import_jac_module(self, node: ast.ModulePath, mod_path: str) -> None:
         """Import a module."""
         self.cur_node = node  # impacts error reporting
         target = import_target_to_relative_path(
-            node.level, node.path_str, os.path.dirname(node.loc.mod_path)
+            level=node.level,
+            target=node.path_str,
+            base_path=os.path.dirname(node.loc.mod_path),
         )
-        return self.import_jac_mod_from_file(target)
+        # If the module is a package (dir)
+        if os.path.isdir(target):
+            self.attach_mod_to_node(node, self.import_jac_mod_from_dir(target))
+            import_node = node.parent_of_type(ast.Import)
+            # And the import is a from import and I am the from module
+            if node == import_node.from_loc:
+                # Import all from items as modules or packages
+                for i in import_node.items.items:
+                    if isinstance(i, ast.ModuleItem):
+                        from_mod_target = import_target_to_relative_path(
+                            level=node.level,
+                            target=node.path_str + "." + i.name.value,
+                            base_path=os.path.dirname(node.loc.mod_path),
+                        )
+                        # If package
+                        if os.path.isdir(from_mod_target):
+                            self.attach_mod_to_node(
+                                i, self.import_jac_mod_from_dir(from_mod_target)
+                            )
+                        # Else module
+                        else:
+                            self.attach_mod_to_node(
+                                i, self.import_jac_mod_from_file(from_mod_target)
+                            )
+        else:
+            self.attach_mod_to_node(node, self.import_jac_mod_from_file(target))
+
+    def import_jac_mod_from_dir(self, target: str) -> ast.Module | None:
+        """Import a module from a directory."""
+        with_init = os.path.join(target, "__init__.jac")
+        if os.path.exists(with_init):
+            return self.import_jac_mod_from_file(with_init)
+        else:
+            return ast.Module(
+                name=target.split(os.path.sep)[-1],
+                source=ast.JacSource("", mod_path=target),
+                doc=None,
+                body=[],
+                is_imported=False,
+                stub_only=True,
+                kid=[ast.EmptyToken()],
+            )
 
     def import_jac_mod_from_file(self, target: str) -> ast.Module | None:
         """Import a module from a file."""

jaclang/compiler/passes/main/pyast_gen_pass.py

@@ -42,6 +42,7 @@ class PyastGenPass(Pass):
         """Initialize pass."""
         self.debuginfo: dict[str, list[str]] = {"jac_mods": []}
         self.already_added: list[str] = []
+        self.method_sigs: list[ast.FuncSignature | ast.EventSignature] = []
         self.preamble: list[ast3.AST] = [
             self.sync(
                 ast3.ImportFrom(
@@ -53,6 +54,13 @@ class PyastGenPass(Pass):
             )
         ]
 
+    def enter_node(self, node: ast.AstNode) -> None:
+        """Enter node."""
+        if node.gen.py_ast:
+            self.prune()
+            return
+        super().enter_node(node)
+
     def exit_node(self, node: ast.AstNode) -> None:
         """Exit node."""
         super().exit_node(node)
@@ -324,7 +332,8 @@ class PyastGenPass(Pass):
         for pbody in node.impl_mod:
             pre_body = [*pre_body, *pbody.body]
         pre_body = [*pre_body, *clean_body]
-        pre_body = [*pre_body, *node.test_mod.body] if node.test_mod else pre_body
+        for pbody in node.test_mod:
+            pre_body = [*pre_body, *pbody.body]
         body = (
             [
                 self.sync(ast3.Expr(value=node.doc.gen.py_ast[0]), jac_node=node.doc),
@@ -658,6 +667,28 @@ class PyastGenPass(Pass):
             )
         ]
 
+    def enter_architype(self, node: ast.Architype) -> None:
+        """Sub objects.
+
+        name: Name,
+        arch_type: Token,
+        access: Optional[SubTag[Token]],
+        base_classes: Optional[SubNodeList[AtomType]],
+        body: Optional[SubNodeList[ArchBlockStmt] | ArchDef],
+        doc: Optional[String],
+        decorators: Optional[SubNodeList[ExprType]],
+        """
+        # Record all signatures that are part of methods
+        for i in (
+            node.body.body.items
+            if isinstance(node.body, ast.ArchDef)
+            else node.body.items if node.body else []
+        ):
+            if isinstance(i, ast.Ability):
+                self.method_sigs.append(i.signature)
+        if isinstance(node.body, ast.AstImplOnlyNode):
+            self.traverse(node.body)
+
     def exit_architype(self, node: ast.Architype) -> None:
         """Sub objects.
 
@@ -821,6 +852,19 @@ class PyastGenPass(Pass):
         decorators: Optional[SubNodeList[ExprType]],
         """
 
+    def enter_enum(self, node: ast.Enum) -> None:
+        """Sub objects.
+
+        name: Name,
+        access: Optional[SubTag[Token]],
+        base_classes: Optional[SubNodeList[AtomType]],
+        body: Optional[SubNodeList[EnumBlockStmt] | EnumDef],
+        doc: Optional[String],
+        decorators: Optional[SubNodeList[ExprType]],
+        """
+        if isinstance(node.body, ast.AstImplOnlyNode):
+            self.traverse(node.body)
+
     def exit_enum(self, node: ast.Enum) -> None:
         """Sub objects.
 
@@ -870,6 +914,23 @@ class PyastGenPass(Pass):
         decorators: Optional[SubNodeList[ExprType]],
         """
 
+    def enter_ability(self, node: ast.Ability) -> None:
+        """Sub objects.
+
+        name_ref: NameType,
+        is_func: bool,
+        is_async: bool,
+        is_static: bool,
+        is_abstract: bool,
+        access: Optional[SubTag[Token]],
+        signature: Optional[FuncSignature | ExprType | EventSignature],
+        body: Optional[SubNodeList[CodeBlockStmt] | AbilityDef | FuncCall],
+        doc: Optional[String],
+        decorators: Optional[SubNodeList[ExprType]],
+        """
+        if isinstance(node.body, ast.AstImplOnlyNode):
+            self.traverse(node.body)
+
     def exit_ability(self, node: ast.Ability) -> None:
         """Sub objects.
 
@@ -1043,7 +1104,7 @@ class PyastGenPass(Pass):
             action = (
                 node.semstr.gen.py_ast[0]
                 if node.semstr
-                else self.sync(ast3.Constant(value=None))
+                else self.sync(ast3.Constant(value=node.name_ref.sym_name))
             )
             return [
                 self.sync(
@@ -1272,7 +1333,7 @@ class PyastGenPass(Pass):
         """
         params = (
             [self.sync(ast3.arg(arg="self", annotation=None))]
-            if node.is_method and not node.is_static
+            if node in self.method_sigs and not node.is_static
             else []
         )
         vararg = None
@@ -1330,7 +1391,7 @@ class PyastGenPass(Pass):
                     posonlyargs=[],
                     args=(
                         [self.sync(ast3.arg(arg="self", annotation=None)), here]
-                        if node.is_method
+                        if node in self.method_sigs
                         else [here]
                     ),
                     kwonlyargs=[],