[stubgenc] Shorten property return types (add necessary imports if needed)

Author: sizmailov

mypy/stubgenc.py

@@ -234,11 +234,14 @@ def strip_or_import(typ: str, module: ModuleType, imports: List[str]) -> str:
     return stripped_type
 
 
-def generate_c_property_stub(name: str, obj: object, output: List[str], readonly: bool) -> None:
+def generate_c_property_stub(name: str, obj: object, output: List[str], readonly: bool,
+                             module: Optional[ModuleType] = None,
+                             imports: Optional[List[str]] = None) -> None:
     """Generate property stub using introspection of 'obj'.
 
     Try to infer type from docstring, append resulting lines to 'output'.
     """
+
     def infer_prop_type(docstr: Optional[str]) -> Optional[str]:
         """Infer property type from docstring or docstring signature."""
         if docstr is not None:
@@ -256,6 +259,9 @@ def infer_prop_type(docstr: Optional[str]) -> Optional[str]:
     if not inferred:
         inferred = 'Any'
 
+    if module is not None and imports is not None:
+        inferred = strip_or_import(inferred, module, imports)
+
     output.append('@property')
     output.append('def {}(self) -> {}: ...'.format(name, inferred))
     if not readonly:
@@ -304,7 +310,8 @@ def generate_c_type_stub(module: ModuleType,
                                          class_sigs=class_sigs)
         elif is_c_property(value):
             done.add(attr)
-            generate_c_property_stub(attr, value, properties, is_c_property_readonly(value))
+            generate_c_property_stub(attr, value, properties, is_c_property_readonly(value),
+                                     module=module, imports=imports)
 
     variables = []
     for attr, value in items:

test-data/stubgen/pybind11_mypy_demo/basics.pyi

@@ -20,27 +20,27 @@ class Point:
     @overload
     def distance_to(*args, **kwargs) -> Any: ...
     @property
-    def angle_unit(self) -> pybind11_mypy_demo.basics.Point.AngleUnit: ...
+    def angle_unit(self) -> Point.AngleUnit: ...
     @angle_unit.setter
-    def angle_unit(self, val: pybind11_mypy_demo.basics.Point.AngleUnit) -> None: ...
+    def angle_unit(self, val: Point.AngleUnit) -> None: ...
     @property
     def length(self) -> float: ...
     @property
-    def length_unit(self) -> pybind11_mypy_demo.basics.Point.LengthUnit: ...
+    def length_unit(self) -> Point.LengthUnit: ...
     @length_unit.setter
-    def length_unit(self, val: pybind11_mypy_demo.basics.Point.LengthUnit) -> None: ...
+    def length_unit(self, val: Point.LengthUnit) -> None: ...
     @property
     def x(self) -> float: ...
     @x.setter
     def x(self, val: float) -> None: ...
     @property
-    def x_axis(self) -> pybind11_mypy_demo.basics.Point: ...
+    def x_axis(self) -> Point: ...
     @property
     def y(self) -> float: ...
     @y.setter
     def y(self, val: float) -> None: ...
     @property
-    def y_axis(self) -> pybind11_mypy_demo.basics.Point: ...
+    def y_axis(self) -> Point: ...
 
 def answer() -> int: ...
 def midpoint(left: float, right: float) -> float: ...