Add unconditional import when type checking

Author: kandluis

pytorch_lightning/__init__.py

@@ -4,6 +4,7 @@
 import os
 import sys
 import time
+import typing
 
 _this_year = time.strftime("%Y")
 __version__ = '1.3.0dev'
@@ -50,18 +51,8 @@
 _PACKAGE_ROOT = os.path.dirname(__file__)
 _PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT)
 
-try:
-    # This variable is injected in the __builtins__ by the build
-    # process. It used to enable importing subpackages of skimage when
-    # the binaries are not built
-    _ = None if __LIGHTNING_SETUP__ else None
-except NameError:
-    __LIGHTNING_SETUP__: bool = False
 
-if __LIGHTNING_SETUP__:  # pragma: no-cover
-    sys.stdout.write(f'Partial import of `{__name__}` during the build process.\n')  # pragma: no-cover
-    # We are not importing the rest of the lightning during the build process, as it may not be compiled yet
-else:
+if typing.TYPE_CHECKING:
     from pytorch_lightning import metrics
     from pytorch_lightning.callbacks import Callback
     from pytorch_lightning.core import LightningDataModule, LightningModule
@@ -76,6 +67,33 @@
         'seed_everything',
         'metrics',
     ]
+else:
+    try:
+        # This variable is injected in the __builtins__ by the build
+        # process. It used to enable importing subpackages of skimage when
+        # the binaries are not built
+        _ = None if __LIGHTNING_SETUP__ else None
+    except NameError:
+        __LIGHTNING_SETUP__: bool = False
+
+    if __LIGHTNING_SETUP__:  # pragma: no-cover
+        sys.stdout.write(f'Partial import of `{__name__}` during the build process.\n')  # pragma: no-cover
+        # We are not importing the rest of the lightning during the build process, as it may not be compiled yet
+    else:
+        from pytorch_lightning import metrics
+        from pytorch_lightning.callbacks import Callback
+        from pytorch_lightning.core import LightningDataModule, LightningModule
+        from pytorch_lightning.trainer import Trainer
+        from pytorch_lightning.utilities.seed import seed_everything
+
+        __all__ = [
+            'Trainer',
+            'LightningDataModule',
+            'LightningModule',
+            'Callback',
+            'seed_everything',
+            'metrics',
+        ]
 
 # for compatibility with namespace packages
 __import__('pkg_resources').declare_namespace(__name__)