skylos 1.0.10__py3-none-any.whl → 2.5.2__py3-none-any.whl

skylos/visitor.py

@@ -1,24 +1,75 @@
 #!/usr/bin/env python3
-import ast,re
+import ast
 from pathlib import Path
+import re
+
+PYTHON_BUILTINS = {
+    "print",
+    "len",
+    "str",
+    "int",
+    "float",
+    "list",
+    "dict",
+    "set",
+    "tuple",
+    "range",
+    "open",
+    "reversed",
+    "super",
+    "object",
+    "type",
+    "enumerate",
+    "zip",
+    "map",
+    "filter",
+    "sorted",
+    "sum",
+    "min",
+    "next",
+    "iter",
+    "bytes",
+    "bytearray",
+    "format",
+    "round",
+    "abs",
+    "complex",
+    "hash",
+    "id",
+    "bool",
+    "callable",
+    "getattr",
+    "max",
+    "all",
+    "any",
+    "setattr",
+    "hasattr",
+    "isinstance",
+    "globals",
+    "locals",
+    "vars",
+    "dir",
+    "property",
+    "classmethod",
+    "staticmethod",
+}
+DYNAMIC_PATTERNS = {"getattr", "globals", "eval", "exec"}
+
+## "🥚" hi :)
 
-PYTHON_BUILTINS={"print","len","str","int","float","list","dict","set","tuple","range","open","super","object","type","enumerate","zip","map","filter","sorted","reversed","sum","min","max","all","any","next","iter","repr","chr","ord","bytes","bytearray","memoryview","format","round","abs","pow","divmod","complex","hash","id","bool","callable","getattr","setattr","delattr","hasattr","isinstance","issubclass","globals","locals","vars","dir","property","classmethod","staticmethod"}
-DYNAMIC_PATTERNS={"getattr","globals","eval","exec"}
 
 class Definition:
-    __slots__ = ('name', 'type', 'filename', 'line', 'simple_name', 'confidence', 'references', 'is_exported', 'in_init')
-    
-    def __init__(self, n, t, f, l):
-        self.name = n
+    def __init__(self, name, t, filename, line):
+        self.name = name
         self.type = t
-        self.filename = f
-        self.line = l
-        self.simple_name = n.split('.')[-1]
+        self.filename = filename
+        self.line = line
+        self.simple_name = name.split(".")[-1]
         self.confidence = 100
         self.references = 0
         self.is_exported = False
-        self.in_init = "__init__.py" in str(f)
-    
+        self.in_init = "__init__.py" in str(filename)
+
     def to_dict(self):
         if self.type == "method" and "." in self.name:
             parts = self.name.split(".")
@@ -28,8 +79,8 @@ class Definition:
                 output_name = self.name
         else:
             output_name = self.simple_name
-            
-        return{
+
+        return {
             "name": output_name,
             "full_name": self.name,
             "simple_name": self.simple_name,
@@ -38,70 +89,125 @@ class Definition:
             "basename": Path(self.filename).name,
             "line": self.line,
             "confidence": self.confidence,
-            "references": self.references
+            "references": self.references,
         }
 
+
 class Visitor(ast.NodeVisitor):
-    def __init__(self,mod,file):
-        self.mod=mod
-        self.file=file
-        self.defs=[]
-        self.refs=[]
-        self.cls=None
-        self.alias={}
-        self.dyn=set()
-        self.exports=set()
+    def __init__(self, mod, file):
+        self.mod = mod
+        self.file = file
+        self.defs = []
+        self.refs = []
+        self.cls = None
+        self.alias = {}
+        self.dyn = set()
+        self.exports = set()
         self.current_function_scope = []
+        self.current_function_params = []
+        self.local_var_maps = []
+        self.in_cst_class = 0
+        self.local_type_maps = []
+        self._dataclass_stack = []
+        self.dataclass_fields = set()
+        self.first_read_lineno = {}
+
+    def add_def(self, name, t, line):
+        found = False
+        for d in self.defs:
+            if d.name == name:
+                found = True
+                break
+        if not found:
+            self.defs.append(Definition(name, t, self.file, line))
 
-    def add_def(self,n,t,l):
-        if n not in{d.name for d in self.defs}:self.defs.append(Definition(n,t,self.file,l))
+    def add_ref(self, name):
+        self.refs.append((name, self.file))
 
-    def add_ref(self,n):self.refs.append((n,self.file))
+    def qual(self, name):
+        if name in self.alias:
+            return self.alias[name]
+        if name in PYTHON_BUILTINS:
+            if self.mod:
+                mod_candidate = f"{self.mod}.{name}"
+            else:
+                mod_candidate = name
+            if any(d.name == mod_candidate for d in self.defs):
+                return mod_candidate
+
+        if self.mod:
+            return f"{self.mod}.{name}"
+        else:
+            return name
 
-    def qual(self,n):
-        if n in self.alias:return self.alias[n]
-        if n in PYTHON_BUILTINS:return n
-        return f"{self.mod}.{n}"if self.mod else n
+    def visit_Global(self, node):
+        if self.current_function_scope and self.local_var_maps:
+            for name in node.names:
+                self.local_var_maps[-1][name] = f"{self.mod}.{name}"
+        return
 
     def visit_annotation(self, node):
         if node is not None:
             if isinstance(node, ast.Constant) and isinstance(node.value, str):
                 self.visit_string_annotation(node.value)
-            elif hasattr(node, 's') and isinstance(node.s, str):
+            elif hasattr(node, "s") and isinstance(node.s, str):
                 self.visit_string_annotation(node.s)
             else:
                 self.visit(node)
 
     def visit_string_annotation(self, annotation_str):
+        if not isinstance(annotation_str, str):
+            return
+
         try:
-            parsed = ast.parse(annotation_str, mode='eval')
+            parsed = ast.parse(annotation_str, mode="eval")
             self.visit(parsed.body)
         except:
-            pass
+            for tok in re.findall(r"[A-Za-z_][A-Za-z0-9_]*", annotation_str):
+                self.add_ref(tok)
 
-    def visit_Import(self,node):
+    def visit_Import(self, node):
         for a in node.names:
-            full=a.name
-            self.alias[a.asname or a.name.split(".")[-1]]=full
-            self.add_def(full,"import",node.lineno)
+            full = a.name
 
-    def visit_ImportFrom(self,node):
-        if node.module is None:return
+            if a.asname:
+                alias_name = a.asname
+                target = full
+            else:
+                head = full.split(".", 1)[0]
+                alias_name = head
+                target = head
+
+            self.alias[alias_name] = target
+            self.add_def(target, "import", node.lineno)
+
+    def visit_ImportFrom(self, node):
+        if node.module is None:
+            return
         for a in node.names:
-            if a.name=="*":continue
-            base=node.module
+            if a.name == "*":
+                continue
+            base = node.module
             if node.level:
-                parts=self.mod.split(".")
-                base=".".join(parts[:-node.level])+(f".{node.module}"if node.module else"")
-            full=f"{base}.{a.name}"
-            self.alias[a.asname or a.name]=full
-            self.add_def(full,"import",node.lineno)
+                parts = self.mod.split(".")
+                base = ".".join(parts[: -node.level]) + (
+                    f".{node.module}" if node.module else ""
+                )
+
+            full = f"{base}.{a.name}"
+            if a.asname:
+                self.alias[a.asname] = full
+                self.add_def(full, "import", node.lineno)
+            else:
+                self.alias[a.name] = full
+                self.add_def(full, "import", node.lineno)
 
     def visit_arguments(self, args):
         for arg in args.args:
             self.visit_annotation(arg.annotation)
         for arg in args.posonlyargs:
             self.visit_annotation(arg.annotation)
+
         for arg in args.kwonlyargs:
             self.visit_annotation(arg.annotation)
         if args.vararg:
@@ -114,55 +220,183 @@ class Visitor(ast.NodeVisitor):
             if default:
                 self.visit(default)
 
-    def visit_FunctionDef(self,node):
-        outer_scope_prefix = '.'.join(self.current_function_scope) + '.' if self.current_function_scope else ''
-        
+    def visit_FunctionDef(self, node):
+        outer_scope_prefix = (
+            ".".join(self.current_function_scope) + "."
+            if self.current_function_scope
+            else ""
+        )
+
         if self.cls:
             name_parts = [self.mod, self.cls, outer_scope_prefix + node.name]
         else:
             name_parts = [self.mod, outer_scope_prefix + node.name]
-        
+
         qualified_name = ".".join(filter(None, name_parts))
 
-        self.add_def(qualified_name,"method"if self.cls else"function",node.lineno)
-        
+        if self.cls:
+            def_type = "method"
+        else:
+            def_type = "function"
+        self.add_def(qualified_name, def_type, node.lineno)
+
         self.current_function_scope.append(node.name)
-        
+        self.local_var_maps.append({})
+        self.local_type_maps.append({})
+
+        old_params = self.current_function_params
+        self.current_function_params = []
+
         for d_node in node.decorator_list:
             self.visit(d_node)
-        
+
+        for arg in node.args.args:
+            param_name = f"{qualified_name}.{arg.arg}"
+            self.add_def(param_name, "parameter", node.lineno)
+            self.current_function_params.append((arg.arg, param_name))
+
         self.visit_arguments(node.args)
         self.visit_annotation(node.returns)
-        
+
         for stmt in node.body:
             self.visit(stmt)
+
         self.current_function_scope.pop()
-        
-    visit_AsyncFunctionDef=visit_FunctionDef
+        self.current_function_params = old_params
+        self.local_var_maps.pop()
+        self.local_type_maps.pop()
+
+    visit_AsyncFunctionDef = visit_FunctionDef
+
+    def visit_ClassDef(self, node):
+        cname = f"{self.mod}.{node.name}"
+        self.add_def(cname, "class", node.lineno)
+
+        is_cst = False
+        is_dc = False
 
-    def visit_ClassDef(self,node):
-        cname=f"{self.mod}.{node.name}"
-        self.add_def(cname,"class",node.lineno)
-        
         for base in node.bases:
+            base_name = ""
+
+            if isinstance(base, ast.Attribute):
+                base_name = base.attr
+
+            elif isinstance(base, ast.Name):
+                base_name = base.id
             self.visit(base)
+
+            if base_name in {"CSTTransformer", "CSTVisitor"}:
+                is_cst = True
+
         for keyword in node.keywords:
             self.visit(keyword.value)
+
         for decorator in node.decorator_list:
+
+            def _is_dc(dec):
+                if isinstance(dec, ast.Call):
+                    target = dec.func
+                else:
+                    target = dec
+
+                if isinstance(target, ast.Name):
+                    return target.id == "dataclass"
+
+                if isinstance(target, ast.Attribute):
+                    return target.attr == "dataclass"
+
+                return False
+
+            if _is_dc(decorator):
+                is_dc = True
             self.visit(decorator)
-        
-        prev=self.cls;self.cls=node.name
-        for b in node.body:self.visit(b)
-        self.cls=prev
+
+        prev = self.cls
+        if is_cst:
+            self.in_cst_class += 1
+
+        self.cls = node.name
+        self._dataclass_stack.append(is_dc)
+        for b in node.body:
+            self.visit(b)
+
+        self.cls = prev
+        self._dataclass_stack.pop()
+
+        if is_cst:
+            self.in_cst_class -= 1
 
     def visit_AnnAssign(self, node):
         self.visit_annotation(node.annotation)
         if node.value:
             self.visit(node.value)
-        self.visit(node.target)
+
+        def _define(t):
+            if isinstance(t, ast.Name):
+                name_simple = t.id
+                scope_parts = [self.mod]
+                if self.cls:
+                    scope_parts.append(self.cls)
+
+                if self.current_function_scope:
+                    scope_parts.extend(self.current_function_scope)
+                prefix = ".".join(filter(None, scope_parts))
+                if prefix:
+                    var_name = f"{prefix}.{name_simple}"
+                else:
+                    var_name = name_simple
+
+                self.add_def(var_name, "variable", t.lineno)
+                if (
+                    self._dataclass_stack
+                    and self._dataclass_stack[-1]
+                    and self.cls
+                    and not self.current_function_scope
+                ):
+                    self.dataclass_fields.add(var_name)
+
+                if self.current_function_scope and self.local_var_maps:
+                    self.local_var_maps[-1][name_simple] = var_name
+
+            elif isinstance(t, (ast.Tuple, ast.List)):
+                for elt in t.elts:
+                    _define(elt)
+
+        _define(node.target)
 
     def visit_AugAssign(self, node):
-        self.visit(node.target)
+        if isinstance(node.target, ast.Name):
+            nm = node.target.id
+            if (
+                self.current_function_scope
+                and self.local_var_maps
+                and self.local_var_maps
+                and nm in self.local_var_maps[-1]
+            ):
+                # self.add_ref(self.local_var_maps[-1][nm])
+                fq = self.local_var_maps[-1][nm]
+                self.add_ref(fq)
+                var_name = fq
+
+            else:
+                self.add_ref(self.qual(nm))
+                scope_parts = [self.mod]
+                if self.cls:
+                    scope_parts.append(self.cls)
+
+                if self.current_function_scope:
+                    scope_parts.extend(self.current_function_scope)
+                prefix = ".".join(filter(None, scope_parts))
+                if prefix:
+                    var_name = f"{prefix}.{nm}"
+                else:
+                    var_name = nm
+
+            self.add_def(var_name, "variable", node.lineno)
+            if self.current_function_scope and self.local_var_maps:
+                self.local_var_maps[-1][nm] = var_name
+        else:
+            self.visit(node.target)
         self.visit(node.value)
 
     def visit_Subscript(self, node):
@@ -178,6 +412,56 @@ class Visitor(ast.NodeVisitor):
             self.visit(node.step)
 
     def visit_Assign(self, node):
+        def process_target_for_def(target_node):
+            if isinstance(target_node, ast.Name):
+                name_simple = target_node.id
+                if (
+                    name_simple == "METADATA_DEPENDENCIES"
+                    and self.cls
+                    and self.in_cst_class > 0
+                ):
+                    return
+                if (
+                    name_simple == "__all__"
+                    and not self.current_function_scope
+                    and not self.cls
+                ):
+                    return
+
+                scope_parts = [self.mod]
+                if self.cls:
+                    scope_parts.append(self.cls)
+                if self.current_function_scope:
+                    scope_parts.extend(self.current_function_scope)
+
+                if (
+                    self.current_function_scope
+                    and self.local_var_maps
+                    and name_simple in self.local_var_maps[-1]
+                ):
+                    var_name = self.local_var_maps[-1][name_simple]
+                else:
+                    prefix = ".".join(filter(None, scope_parts))
+                    if prefix:
+                        var_name = f"{prefix}.{name_simple}"
+                    else:
+                        var_name = name_simple
+
+                self.add_def(var_name, "variable", target_node.lineno)
+                if (
+                    self.current_function_scope
+                    and self.local_var_maps
+                    and name_simple not in self.local_var_maps[-1]
+                ):
+                    self.local_var_maps[-1][name_simple] = var_name
+
+            elif isinstance(target_node, (ast.Tuple, ast.List)):
+                for elt in target_node.elts:
+                    process_target_for_def(elt)
+
+        for t in node.targets:
+            process_target_for_def(t)
+
         for target in node.targets:
             if isinstance(target, ast.Name) and target.id == "__all__":
                 if isinstance(node.value, (ast.List, ast.Tuple)):
@@ -185,47 +469,175 @@ class Visitor(ast.NodeVisitor):
                         value = None
                         if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
                             value = elt.value
-                        elif hasattr(elt, 's') and isinstance(elt.s, str):
+                        elif hasattr(elt, "s") and isinstance(elt.s, str):
                             value = elt.s
-                            
+
                         if value is not None:
-                            full_name = f"{self.mod}.{value}"
-                            self.add_ref(full_name)
+                            if self.mod:
+                                export_name = f"{self.mod}.{value}"
+                            else:
+                                export_name = value
+
+                            self.add_ref(export_name)
                             self.add_ref(value)
+
+        try:
+            if isinstance(node.value, ast.Call):
+                callee = node.value.func
+                fqname = None
+
+                if isinstance(callee, ast.Name):
+                    fqname = self.alias.get(callee.id, self.qual(callee.id))
+
+                elif isinstance(callee, ast.Attribute):
+                    parts = []
+                    cur = callee
+                    while isinstance(cur, ast.Attribute):
+                        parts.append(cur.attr)
+                        cur = cur.value
+                    head = None
+                    if isinstance(cur, ast.Name):
+                        head = self.alias.get(cur.id, self.qual(cur.id))
+                    if head:
+                        fqname = ".".join([head] + list(reversed(parts)))
+
+                if fqname and self.current_function_scope and self.local_type_maps:
+
+                    def _mark_target(t):
+                        if isinstance(t, ast.Name):
+                            self.local_type_maps[-1][t.id] = fqname
+                        elif isinstance(t, (ast.Tuple, ast.List)):
+                            for elt in t.elts:
+                                _mark_target(elt)
+
+                    for t in node.targets:
+                        _mark_target(t)
+        except Exception:
+            pass
+
         self.generic_visit(node)
 
     def visit_Call(self, node):
         self.generic_visit(node)
-        
-        if isinstance(node.func, ast.Name) and node.func.id in ("getattr", "hasattr") and len(node.args) >= 2:
-            if isinstance(node.args[1], ast.Constant) and isinstance(node.args[1].value, str):
+
+        if (
+            isinstance(node.func, ast.Name)
+            and node.func.id in ("getattr", "hasattr")
+            and len(node.args) >= 2
+        ):
+            if isinstance(node.args[1], ast.Constant) and isinstance(
+                node.args[1].value, str
+            ):
                 attr_name = node.args[1].value
                 self.add_ref(attr_name)
-                
+
                 if isinstance(node.args[0], ast.Name):
                     module_name = node.args[0].id
-                    if module_name != "self": 
+                    if module_name != "self":
                         qualified_name = f"{self.qual(module_name)}.{attr_name}"
                         self.add_ref(qualified_name)
-        
+
+            elif isinstance(node.args[0], ast.Name):
+                target_name = node.args[0].id
+                if target_name != "self":
+                    if self.mod:
+                        self.dyn.add(self.mod.split(".")[0])
+                    else:
+                        self.dyn.add("")
+
         elif isinstance(node.func, ast.Name) and node.func.id == "globals":
-            parent = getattr(node, 'parent', None)
-            if (isinstance(parent, ast.Subscript) and 
-                isinstance(parent.slice, ast.Constant) and 
-                isinstance(parent.slice.value, str)):
+            parent = getattr(node, "parent", None)
+            if (
+                isinstance(parent, ast.Subscript)
+                and isinstance(parent.slice, ast.Constant)
+                and isinstance(parent.slice.value, str)
+            ):
                 func_name = parent.slice.value
                 self.add_ref(func_name)
                 self.add_ref(f"{self.mod}.{func_name}")
 
-    def visit_Name(self,node):
-        if isinstance(node.ctx,ast.Load):
-            self.add_ref(self.qual(node.id))
-            if node.id in DYNAMIC_PATTERNS:self.dyn.add(self.mod.split(".")[0])
+        elif isinstance(node.func, ast.Name) and node.func.id in ("eval", "exec"):
+            root_mod = ""
+            if self.mod:
+                root_mod = self.mod.split(".")[0]
+            self.dyn.add(root_mod)
+
+    def visit_Name(self, node):
+        if isinstance(node.ctx, ast.Load):
+            for param_name, param_full_name in self.current_function_params:
+                if node.id == param_name:
+                    self.first_read_lineno.setdefault(param_full_name, node.lineno)
+                    self.add_ref(param_full_name)
+                    return
+
+            if (
+                self.current_function_scope
+                and self.local_var_maps
+                and node.id in self.local_var_maps[-1]
+            ):
+                fq = self.local_var_maps[-1][node.id]
+                self.first_read_lineno.setdefault(fq, node.lineno)
+                self.add_ref(fq)
+                return
 
-    def visit_Attribute(self,node):
+            qualified = self.qual(node.id)
+            self.first_read_lineno.setdefault(qualified, node.lineno)
+            self.add_ref(qualified)
+            if node.id in DYNAMIC_PATTERNS:
+                self.dyn.add(self.mod.split(".")[0])
+
+    def visit_Attribute(self, node):
         self.generic_visit(node)
-        if isinstance(node.ctx,ast.Load)and isinstance(node.value,ast.Name):
-            self.add_ref(f"{self.qual(node.value.id)}.{node.attr}")
+
+        if not isinstance(node.ctx, ast.Load):
+            return
+
+        if isinstance(node.value, ast.Name):
+            base = node.value.id
+
+            param_hit = None
+            for param_name, param_full in self.current_function_params:
+                if base == param_name:
+                    param_hit = (param_name, param_full)
+                    break
+
+            if param_hit:
+                self.add_ref(param_hit[1])
+
+            if self.cls and base in {"self", "cls"}:
+                if self.mod:
+                    owner = f"{self.mod}.{self.cls}"
+                else:
+                    owner = self.cls
+
+                self.add_ref(f"{owner}.{node.attr}")
+                return
+
+            if (
+                self.current_function_scope
+                and self.local_type_maps
+                and self.local_type_maps[-1].get(base)
+            ):
+                self.add_ref(f"{self.local_type_maps[-1][base]}.{node.attr}")
+                return
+
+            self.add_ref(f"{self.qual(base)}.{node.attr}")
+
+    def visit_NamedExpr(self, node):
+        self.visit(node.value)
+        if isinstance(node.target, ast.Name):
+            nm = node.target.id
+            scope_parts = [self.mod]
+            if self.cls:
+                scope_parts.append(self.cls)
+            if self.current_function_scope:
+                scope_parts.extend(self.current_function_scope)
+            prefix = ".".join(filter(None, scope_parts))
+            var_name = f"{prefix}.{nm}" if prefix else nm
+            self.add_def(var_name, "variable", node.lineno)
+            if self.current_function_scope and self.local_var_maps:
+                self.local_var_maps[-1][nm] = var_name
+            self.add_ref(var_name)
 
     def visit_keyword(self, node):
         self.visit(node.value)
@@ -250,4 +662,4 @@ class Visitor(ast.NodeVisitor):
                         self.visit(item)
             elif isinstance(value, ast.AST):
                 value.parent = node
-                self.visit(value)
+                self.visit(value)