stubmaker==0.0.2 release

make_stubs.py

@@ -1,7 +1,7 @@
 import os
 from argparse import ArgumentParser
 
-from stubmaker.builder.ast_builder import ASTBuilder
+from stubmaker.builder.representations_tree_builder import RepresentationsTreeBuilder
 from stubmaker.builder import override_module_import_path, traverse_modules
 from stubmaker.viewers.stub_viewer import StubViewer
 
@@ -32,7 +32,7 @@ def main():
         # Actually creating a file
         print(f'{module_name} -> {dst_path}')
         with open(dst_path, 'w') as stub_flo:
-            builder = ASTBuilder(module_name, module)
+            builder = RepresentationsTreeBuilder(module_name, module)
 
             viewer = StubViewer()
             module_view = viewer.view(builder.module_rep)

setup.py

@@ -13,7 +13,7 @@
     package_dir={PREFIX: 'src'},
     packages=[f'{PREFIX}.{package}' for package in find_packages('src')],
     py_modules=['make_stubs'],
-    version='0.0.1',
+    version='0.0.2',
     description='Tool for generating python stubs',
     long_description=readme,
     long_description_content_type='text/markdown',

src/builder/common.py

@@ -1,3 +1,4 @@
+import inspect
 from typing import Generator, Optional
 
 URL_REGEXP = r'https?://\S+'
@@ -41,9 +42,9 @@ def full_name(self):
 
 class BaseRepresentation:
 
-    def __init__(self, node: Node, ast: 'ASTBuilder'):
+    def __init__(self, node: Node, tree: 'RepresentationsTreeBuilder'):
         self.node = node
-        self.ast = ast
+        self.tree = tree
 
     def __iter__(self):
         raise NotImplementedError
@@ -87,10 +88,10 @@ def __iter__(self):
         raise NotImplementedError
 
     def get_member_rep(self, member_name: str):
-        return self.ast.get_definition(self.node.get_member(member_name))
+        return self.tree.get_definition(self.node.get_member(member_name))
 
 
-class BaseASTBuilder:
+class BaseRepresentationsTreeBuilder:
 
     # Helper methods
 
@@ -130,3 +131,12 @@ def get_literal_for_type_hint(self, obj) -> BaseLiteral:
 
     def get_literal_for_value(self, obj: Node) -> BaseLiteral:
         raise NotImplementedError
+
+
+def get_annotations(obj):
+    if inspect.isclass(obj):
+        annotations = {}
+        for parent in obj.__mro__[::-1]:
+            annotations.update(getattr(parent, '__annotations__', {}))
+        return annotations
+    return getattr(obj, '__annotations__', {})

src/builder/definitions/attribute_annotation_def.py

@@ -1,12 +1,14 @@
-from stubmaker.builder.common import Node, BaseDefinition, BaseASTBuilder
+from stubmaker.builder.common import Node, BaseDefinition, BaseRepresentationsTreeBuilder
 
 
 class AttributeAnnotationDef(BaseDefinition):
     """Represents `name: annotation`"""
 
-    def __init__(self, node: Node, ast: BaseASTBuilder):
-        super().__init__(node, ast)
-        self.annotation = self.ast.get_literal(node)
+    def __init__(self, node: Node, tree: BaseRepresentationsTreeBuilder):
+        super().__init__(node, tree)
+        # we don't want to associate annotation object with name (e.g. TypeVar used in annotation shouldn't be accessed
+        # with this name)
+        self.annotation = self.tree.get_literal(Node(node.namespace, '', node.obj))
 
     @property
     def id(self):

src/builder/definitions/attribute_def.py

@@ -1,12 +1,12 @@
-from stubmaker.builder.common import Node, BaseDefinition, BaseASTBuilder
+from stubmaker.builder.common import Node, BaseDefinition, BaseRepresentationsTreeBuilder
 
 
 class AttributeDef(BaseDefinition):
     """Represents `name = value`"""
 
-    def __init__(self, node: Node, ast: BaseASTBuilder):
-        super().__init__(node, ast)
-        self.value = self.ast.get_literal(node)
+    def __init__(self, node: Node, tree: BaseRepresentationsTreeBuilder):
+        super().__init__(node, tree)
+        self.value = self.tree.get_literal(node)
 
     @property
     def id(self):

src/builder/definitions/class_def.py

@@ -1,22 +1,22 @@
 import inspect
-from typing import get_type_hints
+import typing
 
-from stubmaker.builder.common import Node, BaseASTBuilder, BaseDefinition
+from stubmaker.builder.common import Node, BaseRepresentationsTreeBuilder, BaseDefinition, get_annotations
 
 
 class ClassDef(BaseDefinition):
 
     # TODO:  support properties
 
-    def __init__(self, node: Node, ast: BaseASTBuilder):
-        super().__init__(node, ast)
+    def __init__(self, node: Node, tree: BaseRepresentationsTreeBuilder):
+        super().__init__(node, tree)
 
-        self.docstring = self.ast.get_docstring(self.node)
+        self.docstring = self.tree.get_docstring(self.node)
 
         self.bases = []
         if self.obj.__bases__ != (object,):
             for base in self.node.obj.__bases__:
-                self.bases.append(self.ast.get_literal(Node(self.namespace, None, base)))
+                self.bases.append(self.tree.get_literal(Node(self.namespace, None, base)))
 
         self.members = {}
         for member_name in self.get_public_member_names():
@@ -30,26 +30,28 @@ def __init__(self, node: Node, ast: BaseASTBuilder):
 
             if isinstance(member, staticmethod):
                 node.obj = member.__func__
-                definition = self.ast.get_static_method_definition(node)
+                definition = self.tree.get_static_method_definition(node)
             elif isinstance(member, classmethod):
                 node.obj = member.__func__
-                definition = self.ast.get_class_method_definition(node)
+                definition = self.tree.get_class_method_definition(node)
             elif inspect.isfunction(member):
-                definition = self.ast.get_function_definition(node)
-            elif inspect.isclass(member) and member.__module__ == self.ast.module_name:
-                definition = self.ast.get_class_definition(node)
+                definition = self.tree.get_function_definition(node)
+            elif inspect.isclass(member) and member.__module__ == self.tree.module_name:
+                definition = self.tree.get_class_definition(node)
+            elif isinstance(member, typing.TypeVar):
+                definition = self.tree.get_attribute_definition(node)
             else:
                 continue
-                # definition = self.ast.get_attribute_definition(node)
 
             self.members[member_name] = definition
 
         self.annotations = {}
-        for member_name, annotation in get_type_hints(self.obj).items():
-            self.annotations[member_name] = self.ast.get_attribute_annotation_definition(Node(
+        annotations = get_annotations(self.obj)
+        for member_name, annotation in annotations.items():
+            self.annotations[member_name] = self.tree.get_attribute_annotation_definition(Node(
                 namespace=f'{self.namespace}.{self.name}' if self.namespace else self.name,
                 name=member_name,
-                obj=get_type_hints(self.obj)[member_name],
+                obj=annotations[member_name],
             ))
 
     def __iter__(self):

src/builder/definitions/enum_def.py

@@ -1,16 +1,16 @@
 from enum import Enum
 
-from stubmaker.builder.common import Node, BaseDefinition, BaseASTBuilder
+from stubmaker.builder.common import Node, BaseDefinition, BaseRepresentationsTreeBuilder
 
 
 class EnumDef(BaseDefinition):
 
-    def __init__(self, node: Node, ast: BaseASTBuilder):
-        super().__init__(node, ast)
+    def __init__(self, node: Node, tree: BaseRepresentationsTreeBuilder):
+        super().__init__(node, tree)
         assert issubclass(node.obj, Enum)
-        self.docstring = self.ast.get_docstring(self.node)
-        self.bases = [self.ast.get_literal(Node(self.namespace, None, base)) for base in self.obj.__bases__]
-        self.enum_dict = {e.name: ast.get_literal(Node(self.namespace, None, e.value)) for e in self.obj}
+        self.docstring = self.tree.get_docstring(self.node)
+        self.bases = [self.tree.get_literal(Node(self.namespace, None, base)) for base in self.obj.__bases__]
+        self.enum_dict = {e.name: tree.get_literal(Node(self.namespace, None, e.value)) for e in self.obj}
 
     def __iter__(self):
         yield from self.bases

src/builder/definitions/function_def.py

@@ -1,31 +1,29 @@
 import inspect
 
-from stubmaker.builder.common import BaseDefinition, Node, BaseASTBuilder
+from stubmaker.builder.common import BaseDefinition, Node, BaseRepresentationsTreeBuilder
 
 
 class FunctionDef(BaseDefinition):
 
-    # TODO: support statimethods and classmethods
-
-    def __init__(self, node: Node, ast: BaseASTBuilder):
-        super().__init__(node, ast)
+    def __init__(self, node: Node, tree: BaseRepresentationsTreeBuilder):
+        super().__init__(node, tree)
 
         signature = inspect.signature(self.obj)
 
         params = []
         for param in signature.parameters.values():
             if param.annotation is not inspect.Parameter.empty:
-                param = param.replace(annotation=ast.get_literal(Node(self.namespace, None, param.annotation)))
+                param = param.replace(annotation=tree.get_literal(Node(self.namespace, None, param.annotation)))
             if param.default is not inspect.Parameter.empty:
-                param = param.replace(default=ast.get_literal(Node(self.name, None, param.default)))
+                param = param.replace(default=tree.get_literal(Node(self.name, None, param.default)))
             params.append(param)
 
         return_annotation = signature.return_annotation
         if return_annotation is not inspect.Parameter.empty:
-            return_annotation = ast.get_literal(Node(self.namespace, None, return_annotation))
+            return_annotation = tree.get_literal(Node(self.namespace, None, return_annotation))
 
         self.signature = signature.replace(parameters=params, return_annotation=return_annotation)
-        self.docstring = ast.get_docstring(node)
+        self.docstring = tree.get_docstring(node)
 
     def __iter__(self):
         if self.docstring:

src/builder/definitions/module_def.py

@@ -1,11 +1,11 @@
 import builtins
 import inspect
+import typing
 from collections import defaultdict
-from typing import Any, Dict, get_type_hints
+from typing import Any, Dict
 
-from stubmaker.builder.common import Node, BaseDefinition, BaseASTBuilder
-from stubmaker.builder.literals.reference_literal import ReferenceLiteral
-from stubmaker.builder.literals.type_hint_literal import TypeHintLiteral
+from stubmaker.builder.common import Node, BaseDefinition, BaseRepresentationsTreeBuilder, get_annotations
+from stubmaker.builder.literals import TypeHintLiteral, TypeVarLiteral, ReferenceLiteral
 
 
 def get_type_name(obj):
@@ -19,33 +19,31 @@ def get_type_name(obj):
 
 class ModuleDef(BaseDefinition):
 
-    # TODO: support imported functions
-    # TODO: support imported classes
+    def __init__(self, node: Node, tree: BaseRepresentationsTreeBuilder):
+        super().__init__(node, tree)
 
-    def __init__(self, node: Node, ast: BaseASTBuilder):
-        super().__init__(node, ast)
-
-        self.docstring = self.ast.get_docstring(self.node)
+        self.docstring = self.tree.get_docstring(self.node)
 
         self.members = {}
-        annotations = get_type_hints(self.obj)
+        # we do not use get_type_hints because we do not want to evaluate forward references
+        annotations = get_annotations(self.obj)
 
         public_members = self.get_public_members()
 
         for member_name, member in self.remove_imported_members(public_members).items():
             # check if member is alias
             if get_type_name(member) != member_name:
                 if member_name in annotations:
-                    self.members[member_name] = self.ast.get_attribute_annotation_definition(Node(
+                    self.members[member_name] = self.tree.get_attribute_annotation_definition(Node(
                         namespace=f'{self.namespace}.{self.name}' if self.namespace else self.name,
                         name=member_name,
                         obj=annotations[member_name],
                     ))
                 else:
-                    self.members[member_name] = self.ast.get_attribute_definition(self.node.get_member(member_name))
+                    self.members[member_name] = self.tree.get_attribute_definition(self.node.get_member(member_name))
             else:
                 node = self.node.get_member(member_name)
-                definition = self.ast.get_definition(node)
+                definition = self.tree.get_definition(node)
                 self.members[member_name] = definition
 
     def __iter__(self):
@@ -56,52 +54,57 @@ def __iter__(self):
 
     def _is_import_necessary(self, member):
         if inspect.ismodule(member):
-            return self.ast.module_name != member.__name__
+            return self.tree.module_name != member.__name__
+
+        if isinstance(member, typing.TypeVar):
+            return member.__module__ != self.tree.module_name
 
         member_name = get_type_name(member)
         if not member_name:
             return False
 
-        return self.ast.module_name != member.__module__ and member_name not in builtins.__dict__
+        return self.tree.module_name != member.__module__ and member_name not in builtins.__dict__
 
     def _get_import_from(self, member, member_name):
         guessed_module = inspect.getmodule(member)
         if member_name in guessed_module.__dict__:
             return member.__module__, member_name
 
-        raise ValueError(f"{member_name} can't be imported from from {guessed_module.__name__}")
-
-    def get_imports(self, used_object_ids):
-        imports = set()
-        from_imports = defaultdict(set)
+        raise ValueError(f"{member_name} can't be imported from {guessed_module.__name__}")
 
-        for curr in self.traverse():
-            if curr.id not in used_object_ids:
-                continue
+    def _try_to_add_import_for_object(self, child_repr, imports):
+        if isinstance(child_repr, (TypeHintLiteral, TypeVarLiteral)):
+            if self._is_import_necessary(child_repr.obj):
+                imports.add(child_repr.obj.__module__)
+        elif isinstance(child_repr, ReferenceLiteral):
+            if not self._is_import_necessary(child_repr.obj):
+                return
 
-            if isinstance(curr, TypeHintLiteral):
-                imports.add('typing')
-
-            if isinstance(curr, ReferenceLiteral):
-                if not self._is_import_necessary(curr.obj):
-                    continue
+            # TODO: not all builtins are available from globals. For instance NoneType
+            if inspect.ismodule(child_repr.obj):
+                module_name = child_repr.obj.__name__
+            else:
+                module_name = child_repr.obj.__module__
 
-                # TODO: not all builtins are available from globals. For instance NoneType
-                if inspect.ismodule(curr.obj):
-                    module_name = curr.obj.__name__
+            # hack urllib3 stubs for mypy
+            if module_name and module_name.startswith('urllib3'):
+                if inspect.ismodule(child_repr.obj):
+                    child_repr.obj.__name__ = f'requests.packages.{module_name}'
+                    module_name = child_repr.obj.__name__
                 else:
-                    module_name = curr.obj.__module__
+                    child_repr.obj.__module__ = f'requests.packages.{module_name}'
+                    module_name = child_repr.obj.__module__
 
-                # hack urllib3 stubs for mypy
-                if module_name and module_name.startswith('urllib3'):
-                    if inspect.ismodule(curr.obj):
-                        curr.obj.__name__ = f'requests.packages.{module_name}'
-                        module_name = curr.obj.__name__
-                    else:
-                        curr.obj.__module__ = f'requests.packages.{module_name}'
-                        module_name = curr.obj.__module__
+            imports.add(module_name)
 
-                imports.add(module_name)
+    def get_imports(self, used_object_ids):
+        imports = set()
+        from_imports = defaultdict(set)
+
+        for curr in self.traverse():
+            if curr.id in used_object_ids:
+                for child_repr in curr.traverse():
+                    self._try_to_add_import_for_object(child_repr, imports)
 
         # try to add unused but specified in __all__ dependencies
         for member_name in self.obj.__all__:
@@ -128,7 +131,7 @@ def remove_imported_members(self, members):
                 continue
 
             module_name = getattr(member, '__module__', None)
-            if module_name and module_name != self.ast.module_name and member_name == get_type_name(member):
+            if module_name and module_name != self.tree.module_name and member_name == get_type_name(member):
                 # imported external symbol that is not an alias defined in current module
                 continue
 
@@ -140,7 +143,7 @@ def get_public_members(self) -> Dict[str, Any]:
         public_members = {member_name: member for member_name, member in self.obj.__dict__.items()
                           if not member_name.startswith('__')}
 
-        for member_name in get_type_hints(self.obj):
+        for member_name in get_annotations(self.obj):
             # try to add module level attributes with annotations but without value that are specified in __all__
             # (such attributes can't be retrieved from __dict__)
             if member_name not in public_members and member_name in self.obj.__all__: