Add annotations for ParserModule (#1526)

Author: cdce8p

astroid/_ast.py

@@ -5,9 +5,8 @@
 import ast
 import sys
 import types
-from collections import namedtuple
 from functools import partial
-from typing import Dict, Optional
+from typing import Dict, List, NamedTuple, Optional, Type
 
 from astroid.const import PY38_PLUS, Context
 
@@ -20,23 +19,21 @@
     except ImportError:
         _ast_py3 = None
 
-FunctionType = namedtuple("FunctionType", ["argtypes", "returns"])
 
+class FunctionType(NamedTuple):
+    argtypes: List[ast.expr]
+    returns: ast.expr
 
-class ParserModule(
-    namedtuple(
-        "ParserModule",
-        [
-            "module",
-            "unary_op_classes",
-            "cmp_op_classes",
-            "bool_op_classes",
-            "bin_op_classes",
-            "context_classes",
-        ],
-    )
-):
-    def parse(self, string: str, type_comments=True):
+
+class ParserModule(NamedTuple):
+    module: types.ModuleType
+    unary_op_classes: Dict[Type[ast.unaryop], str]
+    cmp_op_classes: Dict[Type[ast.cmpop], str]
+    bool_op_classes: Dict[Type[ast.boolop], str]
+    bin_op_classes: Dict[Type[ast.operator], str]
+    context_classes: Dict[Type[ast.expr_context], Context]
+
+    def parse(self, string: str, type_comments: bool = True) -> ast.Module:
         if self.module is _ast_py3:
             if PY38_PLUS:
                 parse_func = partial(self.module.parse, type_comments=type_comments)
@@ -58,7 +55,7 @@ def parse_function_type_comment(type_comment: str) -> Optional[FunctionType]:
     return FunctionType(argtypes=func_type.argtypes, returns=func_type.returns)
 
 
-def get_parser_module(type_comments=True) -> ParserModule:
+def get_parser_module(type_comments: bool = True) -> ParserModule:
     parser_module = ast
     if type_comments and _ast_py3:
         parser_module = _ast_py3
@@ -79,11 +76,15 @@ def get_parser_module(type_comments=True) -> ParserModule:
     )
 
 
-def _unary_operators_from_module(module):
+def _unary_operators_from_module(
+    module: types.ModuleType,
+) -> Dict[Type[ast.unaryop], str]:
     return {module.UAdd: "+", module.USub: "-", module.Not: "not", module.Invert: "~"}
 
 
-def _binary_operators_from_module(module):
+def _binary_operators_from_module(
+    module: types.ModuleType,
+) -> Dict[Type[ast.operator], str]:
     binary_operators = {
         module.Add: "+",
         module.BitAnd: "&",
@@ -102,11 +103,15 @@ def _binary_operators_from_module(module):
     return binary_operators
 
 
-def _bool_operators_from_module(module):
+def _bool_operators_from_module(
+    module: types.ModuleType,
+) -> Dict[Type[ast.boolop], str]:
     return {module.And: "and", module.Or: "or"}
 
 
-def _compare_operators_from_module(module):
+def _compare_operators_from_module(
+    module: types.ModuleType,
+) -> Dict[Type[ast.cmpop], str]:
     return {
         module.Eq: "==",
         module.Gt: ">",
@@ -121,7 +126,9 @@ def _compare_operators_from_module(module):
     }
 
 
-def _contexts_from_module(module) -> Dict[ast.expr_context, Context]:
+def _contexts_from_module(
+    module: types.ModuleType,
+) -> Dict[Type[ast.expr_context], Context]:
     return {
         module.Load: Context.Load,
         module.Store: Context.Store,