Fully annotate stubgen (python#2667)

ddfisher · gvanrossum · commit c05b24494476 · 2017-01-09T18:21:32.000-08:00
diff --git a/mypy.ini b/mypy.ini
@@ -2,15 +2,6 @@
 disallow_untyped_defs = True
 
 ; historical exceptions
-[mypy-mypy.stubgen]
-disallow_untyped_defs = False
-
-[mypy-mypy.stubutil]
-disallow_untyped_defs = False
-
-[mypy-mypy.stubgenc]
-disallow_untyped_defs = False
-
 [mypy-mypy.test.testextensions]
 disallow_untyped_defs = False
 
diff --git a/mypy/stubgen.py b/mypy/stubgen.py
@@ -45,7 +45,9 @@
 import sys
 import textwrap
 
-from typing import Any, List, Dict, Tuple, Iterable, Optional, NamedTuple, Set
+from typing import (
+    Any, List, Dict, Tuple, Iterable, Iterator, Optional, NamedTuple, Set, Union, cast
+)
 
 import mypy.build
 import mypy.parse
@@ -341,6 +343,7 @@ def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
 
         for lvalue in o.lvalues:
             if isinstance(lvalue, NameExpr) and self.is_namedtuple(o.rvalue):
+                assert isinstance(o.rvalue, CallExpr)
                 self.process_namedtuple(lvalue, o.rvalue)
                 continue
             if isinstance(lvalue, TupleExpr):
@@ -374,15 +377,15 @@ def is_namedtuple(self, expr: Expression) -> bool:
         return ((isinstance(callee, NameExpr) and callee.name.endswith('namedtuple')) or
                 (isinstance(callee, MemberExpr) and callee.name == 'namedtuple'))
 
-    def process_namedtuple(self, lvalue, rvalue):
+    def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
         self.add_import_line('from collections import namedtuple\n')
         if self._state != EMPTY:
             self.add('\n')
         name = repr(getattr(rvalue.args[0], 'value', '<ERROR>'))
         if isinstance(rvalue.args[1], StrExpr):
             items = repr(rvalue.args[1].value)
         elif isinstance(rvalue.args[1], ListExpr):
-            list_items = rvalue.args[1].items
+            list_items = cast(List[StrExpr], rvalue.args[1].items)
             items = '[%s]' % ', '.join(repr(item.value) for item in list_items)
         else:
             items = '<ERROR>'
@@ -585,7 +588,7 @@ def get_qualified_name(o: Expression) -> str:
         return '<ERROR>'
 
 
-def walk_packages(packages: List[str]):
+def walk_packages(packages: List[str]) -> Iterator[str]:
     for package_name in packages:
         package = __import__(package_name)
         yield package.__name__
diff --git a/mypy/stubgenc.py b/mypy/stubgenc.py
@@ -6,28 +6,34 @@
 import importlib
 import os.path
 import re
-
+from typing import List, Dict, Tuple
+from types import ModuleType
 
 from mypy.stubutil import (
     parse_all_signatures, find_unique_signatures, is_c_module, write_header,
     infer_sig_from_docstring
 )
 
 
-def generate_stub_for_c_module(module_name, target, add_header=True, sigs={}, class_sigs={}):
+def generate_stub_for_c_module(module_name: str,
+                               target: str,
+                               add_header: bool = True,
+                               sigs: Dict[str, str] = {},
+                               class_sigs: Dict[str, str] = {},
+                               ) -> None:
     module = importlib.import_module(module_name)
     assert is_c_module(module), '%s is not a C module' % module_name
     subdir = os.path.dirname(target)
     if subdir and not os.path.isdir(subdir):
         os.makedirs(subdir)
-    functions = []
+    functions = []  # type: List[str]
     done = set()
     items = sorted(module.__dict__.items(), key=lambda x: x[0])
     for name, obj in items:
         if is_c_function(obj):
             generate_c_function_stub(module, name, obj, functions, sigs=sigs)
             done.add(name)
-    types = []
+    types = []  # type: List[str]
     for name, obj in items:
         if name.startswith('__') and name.endswith('__'):
             continue
@@ -62,7 +68,7 @@ def generate_stub_for_c_module(module_name, target, add_header=True, sigs={}, cl
             file.write('%s\n' % line)
 
 
-def add_typing_import(output):
+def add_typing_import(output: List[str]) -> List[str]:
     names = []
     for name in ['Any']:
         if any(re.search(r'\b%s\b' % name, line) for line in output):
@@ -73,27 +79,34 @@ def add_typing_import(output):
         return output[:]
 
 
-def is_c_function(obj):
+def is_c_function(obj: object) -> bool:
     return type(obj) is type(ord)
 
 
-def is_c_method(obj):
+def is_c_method(obj: object) -> bool:
     return type(obj) in (type(str.index),
                          type(str.__add__),
                          type(str.__new__))
 
 
-def is_c_classmethod(obj):
+def is_c_classmethod(obj: object) -> bool:
     type_str = type(obj).__name__
     return type_str == 'classmethod_descriptor'
 
 
-def is_c_type(obj):
+def is_c_type(obj: object) -> bool:
     return type(obj) is type(int)
 
 
-def generate_c_function_stub(module, name, obj, output, self_var=None, sigs={}, class_name=None,
-                             class_sigs={}):
+def generate_c_function_stub(module: ModuleType,
+                             name: str,
+                             obj: object,
+                             output: List[str],
+                             self_var: str = None,
+                             sigs: Dict[str, str] = {},
+                             class_name: str = None,
+                             class_sigs: Dict[str, str] = {},
+                             ) -> None:
     if self_var:
         self_arg = '%s, ' % self_var
     else:
@@ -117,7 +130,13 @@ def generate_c_function_stub(module, name, obj, output, self_var=None, sigs={},
     output.append('def %s(%s%s): ...' % (name, self_arg, sig))
 
 
-def generate_c_type_stub(module, class_name, obj, output, sigs={}, class_sigs={}):
+def generate_c_type_stub(module: ModuleType,
+                         class_name: str,
+                         obj: type,
+                         output: List[str],
+                         sigs: Dict[str, str] = {},
+                         class_sigs: Dict[str, str] = {},
+                         ) -> None:
     items = sorted(obj.__dict__.items(), key=lambda x: method_name_sort_key(x[0]))
     methods = []
     done = set()
@@ -151,7 +170,7 @@ def generate_c_type_stub(module, class_name, obj, output, sigs={}, class_sigs={}
         # TODO: Is this always object?
         del all_bases[-1]
     # Remove base classes of other bases as redundant.
-    bases = []
+    bases = []  # type: List[type]
     for base in all_bases:
         if not any(issubclass(b, base) for b in bases):
             bases.append(base)
@@ -169,15 +188,15 @@ def generate_c_type_stub(module, class_name, obj, output, sigs={}, class_sigs={}
             output.append('    %s' % method)
 
 
-def method_name_sort_key(name):
+def method_name_sort_key(name: str) -> Tuple[int, str]:
     if name in ('__new__', '__init__'):
         return (0, name)
     if name.startswith('__') and name.endswith('__'):
         return (2, name)
     return (1, name)
 
 
-def is_skipped_attribute(attr):
+def is_skipped_attribute(attr: str) -> bool:
     return attr in ('__getattribute__',
                     '__str__',
                     '__repr__',
@@ -187,7 +206,7 @@ def is_skipped_attribute(attr):
                     '__weakref__')  # For pickling
 
 
-def infer_method_sig(name):
+def infer_method_sig(name: str) -> str:
     if name.startswith('__') and name.endswith('__'):
         name = name[2:-2]
         if name in ('hash', 'iter', 'next', 'sizeof', 'copy', 'deepcopy', 'reduce', 'getinitargs',
diff --git a/mypy/stubutil.py b/mypy/stubutil.py
@@ -1,7 +1,8 @@
 import re
 import sys
 
-from typing import Any, Optional, Tuple, Sequence, MutableSequence, List, MutableMapping
+from typing import Any, Optional, Tuple, Sequence, MutableSequence, List, MutableMapping, IO
+from types import ModuleType
 
 
 # Type Alias for Signatures
@@ -85,11 +86,11 @@ def find_unique_signatures(sigs: Sequence[Sig]) -> List[Sig]:
     return sorted(result)
 
 
-def is_c_module(module):
+def is_c_module(module: ModuleType) -> bool:
     return '__file__' not in module.__dict__ or module.__dict__['__file__'].endswith('.so')
 
 
-def write_header(file, module_name, pyversion=(3, 5)):
+def write_header(file: IO[str], module_name: str, pyversion: Tuple[int, int] = (3, 5)) -> None:
     if module_name:
         if pyversion[0] >= 3:
             version = '%d.%d' % (sys.version_info.major,
@@ -102,7 +103,7 @@ def write_header(file, module_name, pyversion=(3, 5)):
         '# NOTE: This dynamically typed stub was automatically generated by stubgen.\n\n')
 
 
-def infer_sig_from_docstring(docstr, name):
+def infer_sig_from_docstring(docstr: str, name: str) -> Optional[str]:
     if not docstr:
         return None
     docstr = docstr.lstrip()
