Typecheck all init files

lightly/api/__init__.py

@@ -2,14 +2,17 @@
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
-from lightly.api import patch as _patch
-from lightly.api.api_workflow_artifacts import ArtifactNotExist
-from lightly.api.api_workflow_client import ApiWorkflowClient
-from lightly.openapi_generated.swagger_client.api_client import (
-    Configuration as _Configuration,
-)
+from typing import TYPE_CHECKING
 
-# Make ApiWorkflowClient and swagger classes picklable.
-_patch.make_swagger_configuration_picklable(
-    configuration_cls=_Configuration,
-)
+if not TYPE_CHECKING:
+    from lightly.api import patch as _patch
+    from lightly.api.api_workflow_artifacts import ArtifactNotExist
+    from lightly.api.api_workflow_client import ApiWorkflowClient
+    from lightly.openapi_generated.swagger_client.api_client import (
+        Configuration as _Configuration,
+    )
+
+    # Make ApiWorkflowClient and swagger classes picklable.
+    _patch.make_swagger_configuration_picklable(
+        configuration_cls=_Configuration,
+    )

lightly/cli/__init__.py

@@ -5,9 +5,11 @@
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
+from typing import TYPE_CHECKING
 
-from lightly.cli.crop_cli import crop_cli
-from lightly.cli.download_cli import download_cli
-from lightly.cli.embed_cli import embed_cli
-from lightly.cli.lightly_cli import lightly_cli
-from lightly.cli.train_cli import train_cli
+if not TYPE_CHECKING:
+    from lightly.cli.crop_cli import crop_cli
+    from lightly.cli.download_cli import download_cli
+    from lightly.cli.embed_cli import embed_cli
+    from lightly.cli.lightly_cli import lightly_cli
+    from lightly.cli.train_cli import train_cli

lightly/data/__init__.py

@@ -2,25 +2,27 @@
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
+from typing import TYPE_CHECKING
 
-from lightly.data._video import (
-    EmptyVideoError,
-    NonIncreasingTimestampError,
-    UnseekableTimestampError,
-    VideoError,
-)
-from lightly.data.collate import (
-    BaseCollateFunction,
-    DINOCollateFunction,
-    ImageCollateFunction,
-    MAECollateFunction,
-    MoCoCollateFunction,
-    MSNCollateFunction,
-    MultiCropCollateFunction,
-    PIRLCollateFunction,
-    SimCLRCollateFunction,
-    SwaVCollateFunction,
-    VICRegLCollateFunction,
-    imagenet_normalize,
-)
-from lightly.data.dataset import LightlyDataset
+if not TYPE_CHECKING:
+    from lightly.data._video import (
+        EmptyVideoError,
+        NonIncreasingTimestampError,
+        UnseekableTimestampError,
+        VideoError,
+    )
+    from lightly.data.collate import (
+        BaseCollateFunction,
+        DINOCollateFunction,
+        ImageCollateFunction,
+        MAECollateFunction,
+        MoCoCollateFunction,
+        MSNCollateFunction,
+        MultiCropCollateFunction,
+        PIRLCollateFunction,
+        SimCLRCollateFunction,
+        SwaVCollateFunction,
+        VICRegLCollateFunction,
+        imagenet_normalize,
+    )
+    from lightly.data.dataset import LightlyDataset

lightly/embedding/__init__.py

@@ -8,5 +8,8 @@
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
 
-from lightly.embedding._base import BaseEmbedding
-from lightly.embedding.embedding import SelfSupervisedEmbedding
+from typing import TYPE_CHECKING
+
+if not TYPE_CHECKING:
+    from lightly.embedding._base import BaseEmbedding
+    from lightly.embedding.embedding import SelfSupervisedEmbedding

lightly/loss/__init__.py

@@ -2,16 +2,18 @@
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
+from typing import TYPE_CHECKING
 
-from lightly.loss.barlow_twins_loss import BarlowTwinsLoss
-from lightly.loss.dcl_loss import DCLLoss, DCLWLoss
-from lightly.loss.dino_loss import DINOLoss
-from lightly.loss.msn_loss import MSNLoss
-from lightly.loss.negative_cosine_similarity import NegativeCosineSimilarity
-from lightly.loss.ntx_ent_loss import NTXentLoss
-from lightly.loss.pmsn_loss import PMSNCustomLoss, PMSNLoss
-from lightly.loss.swav_loss import SwaVLoss
-from lightly.loss.sym_neg_cos_sim_loss import SymNegCosineSimilarityLoss
-from lightly.loss.tico_loss import TiCoLoss
-from lightly.loss.vicreg_loss import VICRegLoss
-from lightly.loss.vicregl_loss import VICRegLLoss
+if not TYPE_CHECKING:
+    from lightly.loss.barlow_twins_loss import BarlowTwinsLoss
+    from lightly.loss.dcl_loss import DCLLoss, DCLWLoss
+    from lightly.loss.dino_loss import DINOLoss
+    from lightly.loss.msn_loss import MSNLoss
+    from lightly.loss.negative_cosine_similarity import NegativeCosineSimilarity
+    from lightly.loss.ntx_ent_loss import NTXentLoss
+    from lightly.loss.pmsn_loss import PMSNCustomLoss, PMSNLoss
+    from lightly.loss.swav_loss import SwaVLoss
+    from lightly.loss.sym_neg_cos_sim_loss import SymNegCosineSimilarityLoss
+    from lightly.loss.tico_loss import TiCoLoss
+    from lightly.loss.vicreg_loss import VICRegLoss
+    from lightly.loss.vicregl_loss import VICRegLLoss

lightly/loss/regularizer/__init__.py

@@ -3,5 +3,7 @@
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
+from typing import TYPE_CHECKING
 
-from lightly.loss.regularizer.co2 import CO2Regularizer
+if not TYPE_CHECKING:
+    from lightly.loss.regularizer.co2 import CO2Regularizer

lightly/models/__init__.py

@@ -17,13 +17,15 @@
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
+from typing import TYPE_CHECKING
 
-from lightly.models import utils
-from lightly.models.barlowtwins import BarlowTwins
-from lightly.models.byol import BYOL
-from lightly.models.moco import MoCo
-from lightly.models.nnclr import NNCLR
-from lightly.models.resnet import ResNetGenerator
-from lightly.models.simclr import SimCLR
-from lightly.models.simsiam import SimSiam
-from lightly.models.zoo import ZOO, checkpoints
+if not TYPE_CHECKING:
+    from lightly.models import utils
+    from lightly.models.barlowtwins import BarlowTwins
+    from lightly.models.byol import BYOL
+    from lightly.models.moco import MoCo
+    from lightly.models.nnclr import NNCLR
+    from lightly.models.resnet import ResNetGenerator
+    from lightly.models.simclr import SimCLR
+    from lightly.models.simsiam import SimSiam
+    from lightly.models.zoo import ZOO, checkpoints

lightly/models/modules/__init__.py

@@ -7,31 +7,33 @@
 
 # Copyright (c) 2021. Lightly AG and its affiliates.
 # All Rights Reserved
+from typing import TYPE_CHECKING
 
-from lightly import _torchvision_vit_available
-from lightly.models.modules.heads import (
-    BarlowTwinsProjectionHead,
-    BYOLPredictionHead,
-    BYOLProjectionHead,
-    DINOProjectionHead,
-    MoCoProjectionHead,
-    NNCLRPredictionHead,
-    NNCLRProjectionHead,
-    SimCLRProjectionHead,
-    SimSiamPredictionHead,
-    SimSiamProjectionHead,
-    SMoGPredictionHead,
-    SMoGProjectionHead,
-    SMoGPrototypes,
-    SwaVProjectionHead,
-    SwaVPrototypes,
-)
-from lightly.models.modules.nn_memory_bank import NNMemoryBankModule
-
-if _torchvision_vit_available:
-    # Requires torchvision >=0.12
-    from lightly.models.modules.masked_autoencoder import (
-        MAEBackbone,
-        MAEDecoder,
-        MAEEncoder,
+if not TYPE_CHECKING:
+    from lightly import _torchvision_vit_available
+    from lightly.models.modules.heads import (
+        BarlowTwinsProjectionHead,
+        BYOLPredictionHead,
+        BYOLProjectionHead,
+        DINOProjectionHead,
+        MoCoProjectionHead,
+        NNCLRPredictionHead,
+        NNCLRProjectionHead,
+        SimCLRProjectionHead,
+        SimSiamPredictionHead,
+        SimSiamProjectionHead,
+        SMoGPredictionHead,
+        SMoGProjectionHead,
+        SMoGPrototypes,
+        SwaVProjectionHead,
+        SwaVPrototypes,
     )
+    from lightly.models.modules.nn_memory_bank import NNMemoryBankModule
+
+    if _torchvision_vit_available:
+        # Requires torchvision >=0.12
+        from lightly.models.modules.masked_autoencoder import (
+            MAEBackbone,
+            MAEDecoder,
+            MAEEncoder,
+        )

lightly/transforms/__init__.py

@@ -7,29 +7,37 @@
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
+from typing import TYPE_CHECKING
 
-from lightly.transforms.byol_transform import (
-    BYOLTransform,
-    BYOLView1Transform,
-    BYOLView2Transform,
-)
-from lightly.transforms.dino_transform import DINOTransform, DINOViewTransform
-from lightly.transforms.fast_siam_transform import FastSiamTransform
-from lightly.transforms.gaussian_blur import GaussianBlur
-from lightly.transforms.jigsaw import Jigsaw
-from lightly.transforms.mae_transform import MAETransform
-from lightly.transforms.moco_transform import MoCoV1Transform, MoCoV2Transform
-from lightly.transforms.msn_transform import MSNTransform, MSNViewTransform
-from lightly.transforms.pirl_transform import PIRLTransform
-from lightly.transforms.rotation import (
-    RandomRotate,
-    RandomRotateDegrees,
-    random_rotation_transform,
-)
-from lightly.transforms.simclr_transform import SimCLRTransform, SimCLRViewTransform
-from lightly.transforms.simsiam_transform import SimSiamTransform, SimSiamViewTransform
-from lightly.transforms.smog_transform import SMoGTransform, SmoGViewTransform
-from lightly.transforms.solarize import RandomSolarization
-from lightly.transforms.swav_transform import SwaVTransform, SwaVViewTransform
-from lightly.transforms.vicreg_transform import VICRegTransform, VICRegViewTransform
-from lightly.transforms.vicregl_transform import VICRegLTransform, VICRegLViewTransform
+if not TYPE_CHECKING:
+    from lightly.transforms.byol_transform import (
+        BYOLTransform,
+        BYOLView1Transform,
+        BYOLView2Transform,
+    )
+    from lightly.transforms.dino_transform import DINOTransform, DINOViewTransform
+    from lightly.transforms.fast_siam_transform import FastSiamTransform
+    from lightly.transforms.gaussian_blur import GaussianBlur
+    from lightly.transforms.jigsaw import Jigsaw
+    from lightly.transforms.mae_transform import MAETransform
+    from lightly.transforms.moco_transform import MoCoV1Transform, MoCoV2Transform
+    from lightly.transforms.msn_transform import MSNTransform, MSNViewTransform
+    from lightly.transforms.pirl_transform import PIRLTransform
+    from lightly.transforms.rotation import (
+        RandomRotate,
+        RandomRotateDegrees,
+        random_rotation_transform,
+    )
+    from lightly.transforms.simclr_transform import SimCLRTransform, SimCLRViewTransform
+    from lightly.transforms.simsiam_transform import (
+        SimSiamTransform,
+        SimSiamViewTransform,
+    )
+    from lightly.transforms.smog_transform import SMoGTransform, SmoGViewTransform
+    from lightly.transforms.solarize import RandomSolarization
+    from lightly.transforms.swav_transform import SwaVTransform, SwaVViewTransform
+    from lightly.transforms.vicreg_transform import VICRegTransform, VICRegViewTransform
+    from lightly.transforms.vicregl_transform import (
+        VICRegLTransform,
+        VICRegLViewTransform,
+    )

lightly/utils/benchmarking/__init__.py

@@ -1,6 +1,11 @@
-from lightly.utils.benchmarking.benchmark_module import BenchmarkModule
-from lightly.utils.benchmarking.knn import knn_predict
-from lightly.utils.benchmarking.knn_classifier import KNNClassifier
-from lightly.utils.benchmarking.linear_classifier import LinearClassifier
-from lightly.utils.benchmarking.metric_callback import MetricCallback
-from lightly.utils.benchmarking.online_linear_classifier import OnlineLinearClassifier
+from typing import TYPE_CHECKING
+
+if not TYPE_CHECKING:
+    from lightly.utils.benchmarking.benchmark_module import BenchmarkModule
+    from lightly.utils.benchmarking.knn import knn_predict
+    from lightly.utils.benchmarking.knn_classifier import KNNClassifier
+    from lightly.utils.benchmarking.linear_classifier import LinearClassifier
+    from lightly.utils.benchmarking.metric_callback import MetricCallback
+    from lightly.utils.benchmarking.online_linear_classifier import (
+        OnlineLinearClassifier,
+    )