2525import re
2626import sys
2727import time
28- from typing import TYPE_CHECKING , Any , Dict , Iterable , NoReturn , cast
28+ from typing import TYPE_CHECKING , Any , Dict , Iterable , NoReturn , Union , cast
2929
3030from mypy .build import BuildSource
3131from mypy .errors import CompileError
4141from mypyc .options import CompilerOptions
4242
4343if TYPE_CHECKING :
44- from distutils .core import Extension
44+ from distutils .core import Extension as _distutils_Extension
45+ from typing_extensions import TypeAlias
46+
47+ from setuptools import Extension as _setuptools_Extension
48+
49+ Extension : TypeAlias = Union [_setuptools_Extension , _distutils_Extension ]
50+
4551
4652try :
4753 # Import setuptools so that it monkey-patch overrides distutils
48- import setuptools # noqa: F401
54+ import setuptools
4955except ImportError :
5056 if sys .version_info >= (3 , 12 ):
5157 # Raise on Python 3.12, since distutils will go away forever
@@ -57,13 +63,16 @@ def get_extension() -> type[Extension]:
5763 # We can work with either setuptools or distutils, and pick setuptools
5864 # if it has been imported.
5965 use_setuptools = "setuptools" in sys .modules
66+ extension_class : type [Extension ]
6067
6168 if not use_setuptools :
62- from distutils .core import Extension
69+ import distutils .core
70+
71+ extension_class = distutils .core .Extension
6372 else :
64- from setuptools import Extension
73+ extension_class = setuptools . Extension
6574
66- return Extension
75+ return extension_class
6776
6877
6978def setup_mypycify_vars () -> None :
0 commit comments