Add support for stub modules semantics

beniget/beniget.py

@@ -290,14 +290,19 @@ class DefUseChains(ast.NodeVisitor):
     One instance of DefUseChains is only suitable to analyse one AST Module in it's lifecycle.
     """
 
-    def __init__(self, filename=None):
+    def __init__(self, filename=None, future_annotations=False, is_stub=False):
         """
             - filename: str, included in error messages if specified
+            - future_annotations: bool, PEP 563 mode 
+            - is_stub: bool, stub module semantics mode, implies future_annotations=True.
+                When the module is a stub, there is no need for quoting to do a forward reference 
+                inside a type alias, typevar bound argument or a classdef base expression.
         """
         self.chains = {}
         self.locals = defaultdict(list)
 
         self.filename = filename
+        self.is_stub = is_stub
 
         # deep copy of builtins, to remain reentrant
         self._builtins = {k: Def(v) for k, v in Builtins.items()}
@@ -339,7 +344,7 @@ def __init__(self, filename=None):
 
         # attributes set in visit_Module
         self.module = None
-        self.future_annotations = False
+        self.future_annotations = is_stub or future_annotations
 
     #
     ## helpers
@@ -720,10 +725,21 @@ def visit_ClassDef(self, node):
         dnode = self.chains.setdefault(node, Def(node))
         self.locals[self._scopes[-1]].append(dnode)
 
-        for base in node.bases:
-            self.visit(base).add_user(dnode)
-        for keyword in node.keywords:
-            self.visit(keyword.value).add_user(dnode)
+        if self.is_stub:
+            # special treatment for base classes in stub modules
+            # so they can contain forward-references.
+            currentscopes = list(self._scopes)
+            for base in node.bases:
+                self._defered_annotations[-1].append((
+                    base, currentscopes, lambda dbase: dbase.add_user(dnode)))
+            for keyword in node.keywords:
+                self._defered_annotations[-1].append((
+                    keyword.value, currentscopes, lambda dkeyword: dkeyword.add_user(dnode)))
+        else:
+            for base in node.bases:
+                self.visit(base).add_user(dnode)
+            for keyword in node.keywords:
+                self.visit(keyword.value).add_user(dnode)
         for decorator in node.decorator_list:
             self.visit(decorator).add_user(dnode)
 
@@ -757,20 +773,31 @@ def visit_Assign(self, node):
         self.visit(node.value)
         for target in node.targets:
             self.visit(target)
-
+    
     def visit_AnnAssign(self, node):
-        if node.value:
+        visit_value = True
+        if (self.is_stub and node.value and _is_typing_name(
+                node.annotation, 'TypeAlias')):
+            # support for PEP 613 - Explicit Type Aliases
+            # BUT an untyped global expression 'x=int' will NOT be considered a type alias.
+            visit_value = False
+            self._defered_annotations[-1].append(
+                (node.value, list(self._scopes),
+                lambda dvalue:dvalue.add_user(dtarget)))
+        elif node.value:
             dvalue = self.visit(node.value)
+
         if not self.future_annotations:
             dannotation = self.visit(node.annotation)
         else:
             self._defered_annotations[-1].append(
                 (node.annotation, list(self._scopes),
                 lambda d:dtarget.add_user(d)))
+
         dtarget = self.visit(node.target)
         if not self.future_annotations:
             dtarget.add_user(dannotation)
-        if node.value:
+        if node.value and visit_value:
             dvalue.add_user(dtarget)
 
     def visit_AugAssign(self, node):
@@ -1104,10 +1131,23 @@ def visit_Compare(self, node):
     def visit_Call(self, node):
         dnode = self.chains.setdefault(node, Def(node))
         self.visit(node.func).add_user(dnode)
-        for arg in node.args:
-            self.visit(arg).add_user(dnode)
-        for kw in node.keywords:
-            self.visit(kw.value).add_user(dnode)
+        if self.is_stub and _is_typing_name(node.func, 'TypeVar'):
+            # In stubs, constraints and bound argument 
+            # of TypeVar() can be forward references.
+            current_scopes = list(self._scopes)
+            for arg in node.args:
+                self._defered_annotations[-1].append(
+                    (arg, current_scopes,
+                    lambda darg:darg.add_user(dnode)))
+            for kw in node.keywords:
+                self._defered_annotations[-1].append(
+                    (kw.value, current_scopes,
+                    lambda dkw:dkw.add_user(dnode)))
+        else:
+            for arg in node.args:
+                self.visit(arg).add_user(dnode)
+            for kw in node.keywords:
+                self.visit(kw.value).add_user(dnode)
         return dnode
 
     visit_Repr = visit_Await
@@ -1263,6 +1303,23 @@ def _iter_arguments(args):
     if args.kwarg:
         yield args.kwarg
 
+def _is_name(expr, name, modules=()):
+    """
+    Returns True if the expression matches:
+     - Name(id=<name>) or
+     - Attribue(value=Name(id=<one of modules>), attr=<name>)
+    """
+    if isinstance(expr, ast.Name):
+        return expr.id==name
+    if isinstance(expr, ast.Attribute):
+        if isinstance(expr.value, ast.Name):
+            if expr.value.id in modules:
+                return expr.attr==name
+    return False
+
+def _is_typing_name(expr, name):
+    return _is_name(expr, name, {'typing', 'typing_extensions', 't'})
+
 def lookup_annotation_name_defs(name, heads, locals_map):
     r"""
     Simple identifier -> defs resolving.

tests/test_chains.py

@@ -24,7 +24,7 @@ def captured_output():
 
 
 class TestDefUseChains(TestCase):
-    def checkChains(self, code, ref, strict=True):
+    def checkChains(self, code, ref, strict=True, is_stub=False):
         class StrictDefUseChains(beniget.DefUseChains):
             def unbound_identifier(self, name, node):
                 raise RuntimeError(
@@ -35,9 +35,9 @@ def unbound_identifier(self, name, node):
 
         node = ast.parse(code)
         if strict:
-            c = StrictDefUseChains()
+            c = StrictDefUseChains(is_stub=is_stub)
         else:
-            c = beniget.DefUseChains()
+            c = beniget.DefUseChains(is_stub=is_stub)
         c.visit(node)
         self.assertEqual(c.dump_chains(node), ref)
         return node, c
@@ -1129,6 +1129,77 @@ class A:
                  'A -> (A -> (), A -> ())'], # good
                 strict=False
             )
+
+    @skipIf(sys.version_info.major < 3, "Python 3 semantics")
+    def test_stubs_generic_base_forward_ref(self):
+        code = '''
+Thing = object
+class _ScandirIterator(str, int, Thing[_ScandirIterator[F]], object):
+    ...
+class C(F, k=H):
+    ...
+F = H = object
+'''
+        self.checkChains(
+                code, 
+                ['Thing -> (Thing -> (Subscript -> (_ScandirIterator -> (_ScandirIterator -> (Subscript -> ((#2)))))))',
+                 '_ScandirIterator -> (_ScandirIterator -> (Subscript -> (Subscript -> ((#0)))))',
+                 'C -> ()',
+                 'F -> (F -> (Subscript -> (Subscript -> (_ScandirIterator -> (_ScandirIterator -> ((#2)))))), F -> (C -> ()))',
+                 'H -> (H -> (C -> ()))'],
+                is_stub=True
+            )
+
+        self.checkChains(
+                code, 
+                ['Thing -> (Thing -> (Subscript -> (_ScandirIterator -> ())))',
+                 '_ScandirIterator -> ()', 
+                 'C -> ()',
+                 'F -> ()',
+                 'H -> ()'],
+                strict=False,
+            )
+
+    @skipIf(sys.version_info.major < 3, "Python 3 semantics")
+    def test_stubs_forward_ref(self):
+        code = '''
+from typing import TypeAlias
+LiteralValue: TypeAlias = list[LiteralValue]|object
+'''
+        self.checkChains(
+                code, 
+                ['TypeAlias -> (TypeAlias -> ())',
+                 'LiteralValue -> (LiteralValue -> (Subscript -> (BinOp -> ((#0)))), TypeAlias -> ())'],
+                is_stub=True
+            )
+        self.checkChains(
+                code, 
+                ['TypeAlias -> (TypeAlias -> ())',
+                 'LiteralValue -> (TypeAlias -> ())'],
+                strict=False,
+            )
+
+    @skipIf(sys.version_info.major < 3, "Python 3 semantics")
+    def test_stubs_typevar_forward_ref(self):
+        code = '''
+from typing import TypeVar
+AnyStr = TypeVar('AnyStr', F, bound=ast.AST)
+import ast
+F = object
+'''
+        self.checkChains(
+                code, 
+                ['TypeVar -> (TypeVar -> (Call -> ()))',
+                 'AnyStr -> ()',
+                 'ast -> (ast -> (Attribute -> (Call -> ())))',
+                 'F -> (F -> (Call -> ()))'],
+                is_stub=True
+            )
+        self.checkChains(
+                code, 
+                ['TypeVar -> (TypeVar -> (Call -> ()))', 'AnyStr -> ()', 'ast -> ()', 'F -> ()'],
+                strict=False,
+            )
 
 class TestUseDefChains(TestCase):
     def checkChains(self, code, ref):