Simplify name adding API

Author: chadrik

mypy/stubgen.py

@@ -516,21 +516,20 @@ def _get_func_return(self, o: FuncDef, ctx: FunctionContext) -> str | None:
         if retname is not None:
             return retname
         if has_yield_expression(o) or has_yield_from_expression(o):
-            self.add_typing_import("Generator")
             yield_name = "None"
             send_name = "None"
             return_name = "None"
             if has_yield_from_expression(o):
-                yield_name = send_name = self.add_incomplete()
+                yield_name = send_name = self.add_name("_typeshed.Incomplete")
             else:
                 for expr, in_assignment in all_yield_expressions(o):
                     if expr.expr is not None and not is_none_expr(expr.expr):
-                        yield_name = self.add_incomplete()
+                        yield_name = self.add_name("_typeshed.Incomplete")
                     if in_assignment:
-                        send_name = self.add_incomplete()
+                        send_name = self.add_name("_typeshed.Incomplete")
             if has_return_statement(o):
-                return_name = self.add_incomplete()
-            generator_name = self.add_typing_import("Generator")
+                return_name = self.add_name("_typeshed.Incomplete")
+            generator_name = self.add_name("collections.abc.Generator")
             return f"{generator_name}[{yield_name}, {send_name}, {return_name}]"
         if not has_return_statement(o) and o.abstract_status == NOT_ABSTRACT:
             return "None"
@@ -731,18 +730,17 @@ def get_base_types(self, cdef: ClassDef) -> list[str]:
                     typename = base.args[0].value
                     if nt_fields is not None:
                         fields_str = ", ".join(f"({f!r}, {t})" for f, t in nt_fields)
-                        namedtuple_name = self.add_typing_import("NamedTuple")
+                        namedtuple_name = self.add_name("typing.NamedTuple")
                         base_types.append(f"{namedtuple_name}({typename!r}, [{fields_str}])")
-                        self.add_typing_import("NamedTuple")
                     else:
                         # Invalid namedtuple() call, cannot determine fields
-                        base_types.append(self.add_incomplete())
+                        base_types.append(self.add_name("_typeshed.Incomplete"))
                 elif self.is_typed_namedtuple(base):
                     base_types.append(base.accept(p))
                 else:
                     # At this point, we don't know what the base class is, so we
                     # just use Incomplete as the base class.
-                    base_types.append(self.add_incomplete())
+                    base_types.append(self.add_name("_typeshed.Incomplete"))
         for name, value in cdef.keywords.items():
             if name == "metaclass":
                 continue  # handled separately
@@ -823,7 +821,7 @@ def _get_namedtuple_fields(self, call: CallExpr) -> list[tuple[str, str]] | None
             else:
                 return None  # Invalid namedtuple fields type
             if field_names:
-                incomplete = self.add_incomplete()
+                incomplete = self.add_name("_typeshed.Incomplete")
                 return [(field_name, incomplete) for field_name in field_names]
             else:
                 return []
@@ -857,7 +855,7 @@ def process_namedtuple(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
         if fields is None:
             self.annotate_as_incomplete(lvalue)
             return
-        bases = self.add_typing_import("NamedTuple")
+        bases = self.add_name("typing.NamedTuple")
         # TODO: Add support for generic NamedTuples. Requires `Generic` as base class.
         class_def = f"{self._indent}class {lvalue.name}({bases}):"
         if len(fields) == 0:
@@ -907,14 +905,13 @@ def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
                     total = arg
                 else:
                     items.append((arg_name, arg))
-        self.add_typing_import("TypedDict")
         p = AliasPrinter(self)
         if any(not key.isidentifier() or keyword.iskeyword(key) for key, _ in items):
             # Keep the call syntax if there are non-identifier or reserved keyword keys.
             self.add(f"{self._indent}{lvalue.name} = {rvalue.accept(p)}\n")
             self._state = VAR
         else:
-            bases = self.add_typing_import("TypedDict")
+            bases = self.add_name("typing_extensions.TypedDict")
             # TODO: Add support for generic TypedDicts. Requires `Generic` as base class.
             if total is not None:
                 bases += f", total={total.accept(p)}"
@@ -931,7 +928,7 @@ def process_typeddict(self, lvalue: NameExpr, rvalue: CallExpr) -> None:
                 self._state = CLASS
 
     def annotate_as_incomplete(self, lvalue: NameExpr) -> None:
-        incomplete = self.add_incomplete()
+        incomplete = self.add_name("_typeshed.Incomplete")
         self.add(f"{self._indent}{lvalue.name}: {incomplete}\n")
         self._state = VAR
 
@@ -1134,10 +1131,10 @@ def get_str_type_of_node(
         if isinstance(rvalue, NameExpr) and rvalue.name in ("True", "False"):
             return "bool"
         if can_infer_optional and isinstance(rvalue, NameExpr) and rvalue.name == "None":
-            incomplete = self.add_incomplete()
+            incomplete = self.add_name("_typeshed.Incomplete")
             return f"{incomplete} | None"
         if can_be_any:
-            return self.add_incomplete()
+            return self.add_name("_typeshed.Incomplete")
         else:
             return ""
 

mypy/stubgenc.py

@@ -288,7 +288,7 @@ def get_annotation(key: str) -> str | None:
                     if argtype == "None":
                         # None is not a useful annotation, but we can infer that the arg
                         # is optional
-                        incomplete = self.add_incomplete()
+                        incomplete = self.add_name("_typeshed.Incomplete")
                         argtype = f"{incomplete} | None"
                 arglist.append(ArgSig(arg, argtype, default=True))
             else:
@@ -459,9 +459,9 @@ def get_type_annotation(self, obj: object) -> str:
         elif inspect.isclass(obj):
             return "type[{}]".format(self.get_type_fullname(obj))
         elif isinstance(obj, FunctionType):
-            return self.add_typing_import("Callable")
+            return self.add_name("typing.Callable")
         elif isinstance(obj, ModuleType):
-            return self.add_obj_import("types", "ModuleType")
+            return self.add_name("types.ModuleType", require=False)
         else:
             return self.get_type_fullname(type(obj))
 
@@ -556,10 +556,10 @@ def generate_function_stub(
 
         decorators = []
         if len(inferred) > 1:
-            decorators.append("@{}".format(self.add_typing_import("overload")))
+            decorators.append("@{}".format(self.add_name("typing.overload")))
 
         if ctx.is_abstract:
-            decorators.append("@{}".format(self.add_abc_import("abstractmethod")))
+            decorators.append("@{}".format(self.add_name("collections.abc.abstractmethod")))
 
         if class_info is not None:
             if self.is_staticmethod(class_info, name, obj):
@@ -646,10 +646,10 @@ def generate_property_stub(
             inferred_type = self.strip_or_import(inferred_type)
 
         if static:
-            classvar = self.add_typing_import("ClassVar")
+            classvar = self.add_name("typing.ClassVar")
             trailing_comment = "  # read-only" if readonly else ""
             if inferred_type is None:
-                inferred_type = self.add_incomplete()
+                inferred_type = self.add_name("_typeshed.Incomplete")
 
             static_properties.append(
                 f"{self._indent}{name}: {classvar}[{inferred_type}] = ...{trailing_comment}"
@@ -661,7 +661,7 @@ def generate_property_stub(
                 ro_properties.append(self._indent + sig.format_sig())
             else:
                 if inferred_type is None:
-                    inferred_type = self.add_incomplete()
+                    inferred_type = self.add_name("_typeshed.Incomplete")
 
                 rw_properties.append(f"{self._indent}{name}: {inferred_type}")
 
@@ -755,7 +755,7 @@ def generate_class_stub(self, class_name: str, cls: type, output: list[str]) ->
                 # special case for __hash__
                 continue
             prop_type_name = self.strip_or_import(self.get_type_annotation(value))
-            classvar = self.add_typing_import("ClassVar")
+            classvar = self.add_name("typing.ClassVar")
             static_properties.append(f"{self._indent}{attr}: {classvar}[{prop_type_name}] = ...")
 
         self.dedent()

mypy/stubutil.py

@@ -552,32 +552,16 @@ def refers_to_fullname(self, name: str, fullname: str | tuple[str, ...]) -> bool
             name == short or self.import_tracker.reverse_alias.get(name) == short
         )
 
-    def add_obj_import(self, module: str, name: str, require: bool = False) -> str:
-        """Add a name to be imported.
+    def add_name(self, fullname: str, require: bool = True) -> str:
+        """Add a name to be imported and return the name reference.
 
         The import will be internal to the stub (i.e don't reexport).
         """
+        module, name = fullname.rsplit(".", 1)
         alias = "_" + name if name in self.defined_names else None
         self.import_tracker.add_import_from(module, [(name, alias)], require=require)
         return alias or name
 
-    def add_typing_import(self, name: str, require: bool = True) -> str:
-        """Add a name to be imported from typing, unless it's imported already.
-
-        The import will be internal to the stub (i.e don't reexport).
-        """
-        return self.add_obj_import("typing", name, require=require)
-
-    def add_incomplete(self, require: bool = True) -> str:
-        return self.add_obj_import("_typeshed", "Incomplete", require=require)
-
-    def add_abc_import(self, name: str, require: bool = True) -> str:
-        """Add a name to be imported from collections.abc, unless it's imported already.
-
-        The import will be internal to the stub.
-        """
-        return self.add_obj_import("collections.abc", name, require=require)
-
     def add_import_line(self, line: str) -> None:
         """Add a line of text to the import section, unless it's already there."""
         if line not in self._import_lines:
@@ -647,7 +631,7 @@ def set_defined_names(self, defined_names: set[str]) -> None:
             for t in imports:
                 # require=False means that the import won't be added unless require_name() is called
                 # for the object during generation.
-                self.add_obj_import(pkg, t, require=False)
+                self.add_name(f"{pkg}.{t}", require=False)
 
     def get_signatures(
         self,