diff --git a/src/stubgen.py b/src/stubgen.py index b7843768..c4ff32ae 100755 --- a/src/stubgen.py +++ b/src/stubgen.py @@ -479,7 +479,10 @@ def put_type(self, tp: NbType, name: Optional[str]): if same_module: # This is an alias of a type in the same module or same top-level module - alias_tp = self.import_object("typing", "TypeAlias") + if sys.version_info >= (3, 10, 0): + alias_tp = self.import_object("typing", "TypeAlias") + else: + alias_tp = self.import_object("typing_extensions", "TypeAlias") self.write_ln(f"{name}: {alias_tp} = {tp.__qualname__}\n") elif self.include_external_imports or (same_toplevel_module and self.include_internal_imports): # Import from a different module @@ -581,7 +584,10 @@ def put_value(self, value: object, name: str, parent: Optional[object] = None, a if self.is_type_var(tp): types = "" elif typing.get_origin(value): - types = ": " + self.import_object("typing", "TypeAlias") + if sys.version_info >= (3, 10, 0): + types = ": " + self.import_object("typing", "TypeAlias") + else: + types = ": " + self.import_object("typing_extensions", "TypeAlias") else: types = f": {self.type_str(tp)}"