improve name handling in stub generation

This commit is contained in:
sneakers-the-rat 2024-05-17 16:33:37 -07:00
parent 7fee64ada6
commit 3a794a57c8
Signed by untrusted user who does not match committer: jonny
GPG key ID: 6DCB96EF1E4D232D

View file

@ -23,7 +23,7 @@ def generate_ndarray_stub() -> str:
# Create import statements, saving aliased name of type if needed # Create import statements, saving aliased name of type if needed
if arr.__module__.startswith("numpydantic") or arr.__module__ == "typing": if arr.__module__.startswith("numpydantic") or arr.__module__ == "typing":
type_name = arr.__name__ type_name = str(arr) if arr.__module__ == "typing" else arr.__name__
import_strings.append(f"from {arr.__module__} import {arr.__name__}") import_strings.append(f"from {arr.__module__} import {arr.__name__}")
else: else:
# since other packages could use the same name for an imported object # since other packages could use the same name for an imported object
@ -36,10 +36,7 @@ def generate_ndarray_stub() -> str:
f"from {arr.__module__} import {arr.__name__} " f"as {type_name}" f"from {arr.__module__} import {arr.__name__} " f"as {type_name}"
) )
if arr.__module__ != "typing":
type_names.append(type_name) type_names.append(type_name)
else:
type_names.append(str(arr))
import_strings.extend(_BUILTIN_IMPORTS) import_strings.extend(_BUILTIN_IMPORTS)
import_string = "\n".join(import_strings) import_string = "\n".join(import_strings)