use ModuleNotFoundError instead of ImportError (#9867)

Co-authored-by: Adrian Wälchli <[email protected]>

CHANGELOG.md

@@ -177,6 +177,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Module imports are now catching `ModuleNotFoundError` instead of `ImportError` ([#9867](https://github.com/PyTorchLightning/pytorch-lightning/pull/9867))
+
 - `pytorch_lightning.loggers.neptune.NeptuneLogger` is now consistent with new [neptune-client](https://github.com/neptune-ai/neptune-client) API ([#6867](https://github.com/PyTorchLightning/pytorch-lightning/pull/6867)).
 
   Old [neptune-client](https://github.com/neptune-ai/neptune-client) API is supported by `NeptuneClient` from [neptune-contrib](https://github.com/neptune-ai/neptune-contrib) repo.

pytorch_lightning/callbacks/progress/rich_progress.py

@@ -198,7 +198,7 @@ class RichProgressBar(ProgressBarBase):
         theme: Contains styles used to stylize the progress bar.
 
     Raises:
-        ImportError:
+        ModuleNotFoundError:
             If required `rich` package is not installed on the device.
     """
 
@@ -208,7 +208,7 @@ def __init__(
         theme: RichProgressBarTheme = RichProgressBarTheme(),
     ) -> None:
         if not _RICH_AVAILABLE:
-            raise ImportError(
+            raise ModuleNotFoundError(
                 "`RichProgressBar` requires `rich` to be installed. Install it by running `pip install -U rich`."
             )
         super().__init__()

pytorch_lightning/callbacks/rich_model_summary.py

@@ -54,13 +54,13 @@ class RichModelSummary(ModelSummary):
             layer summary off.
 
     Raises:
-        ImportError:
+        ModuleNotFoundError:
             If required `rich` package is not installed on the device.
     """
 
     def __init__(self, max_depth: int = 1) -> None:
         if not _RICH_AVAILABLE:
-            raise ImportError(
+            raise ModuleNotFoundError(
                 "`RichProgressBar` requires `rich` to be installed. Install it by running `pip install -U rich`."
             )
         super().__init__(max_depth)

pytorch_lightning/loggers/__init__.py

@@ -27,7 +27,7 @@
 
 if _COMET_AVAILABLE:
     __all__.append("CometLogger")
-    # needed to prevent ImportError and duplicated logs.
+    # needed to prevent ModuleNotFoundError and duplicated logs.
     environ["COMET_DISABLE_AUTO_LOGGING"] = "1"
 
 if _MLFLOW_AVAILABLE:

pytorch_lightning/loggers/comet.py

@@ -40,7 +40,7 @@
 
     try:
         from comet_ml.api import API
-    except ImportError:  # pragma: no-cover
+    except ModuleNotFoundError:  # pragma: no-cover
         # For more information, see: https://www.comet.ml/docs/python-sdk/releases/#release-300
         from comet_ml.papi import API  # pragma: no-cover
 else:
@@ -120,7 +120,7 @@ class CometLogger(LightningLoggerBase):
             :class:`CometExperiment` can be passed as keyword arguments in this logger.
 
     Raises:
-        ImportError:
+        ModuleNotFoundError:
             If required Comet package is not installed on the device.
         MisconfigurationException:
             If neither ``api_key`` nor ``save_dir`` are passed as arguments.
@@ -141,7 +141,7 @@ def __init__(
         **kwargs,
     ):
         if comet_ml is None:
-            raise ImportError(
+            raise ModuleNotFoundError(
                 "You want to use `comet_ml` logger which is not installed yet, install it with `pip install comet-ml`."
             )
         super().__init__()

pytorch_lightning/loggers/mlflow.py

@@ -33,7 +33,7 @@
     from mlflow.tracking import context, MlflowClient
     from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME
 # todo: there seems to be still some remaining import error with Conda env
-except ImportError:
+except ModuleNotFoundError:
     _MLFLOW_AVAILABLE = False
     mlflow, MlflowClient, context = None, None, None
     MLFLOW_RUN_NAME = "mlflow.runName"
@@ -100,7 +100,7 @@ def any_lightning_module_function_or_hook(self):
             default.
 
     Raises:
-        ImportError:
+        ModuleNotFoundError:
             If required MLFlow package is not installed on the device.
     """
 
@@ -117,7 +117,7 @@ def __init__(
         artifact_location: Optional[str] = None,
     ):
         if mlflow is None:
-            raise ImportError(
+            raise ModuleNotFoundError(
                 "You want to use `mlflow` logger which is not installed yet, install it with `pip install mlflow`."
             )
         super().__init__()

pytorch_lightning/loggers/neptune.py

@@ -42,7 +42,7 @@
         from neptune.new.exceptions import NeptuneLegacyProjectException, NeptuneOfflineModeFetchException
         from neptune.new.run import Run
         from neptune.new.types import File as NeptuneFile
-    except ImportError:
+    except ModuleNotFoundError:
         import neptune
         from neptune.exceptions import NeptuneLegacyProjectException
         from neptune.run import Run
@@ -243,7 +243,7 @@ def any_lightning_module_function_or_hook(self):
             used when run is created.
 
     Raises:
-        ImportError:
+        ModuleNotFoundError:
             If required Neptune package in version >=0.9 is not installed on the device.
         TypeError:
             If configured project has not been migrated to new structure yet.

pytorch_lightning/loggers/test_tube.py

@@ -83,7 +83,7 @@ def any_lightning_module_function_or_hook(self):
         prefix: A string to put at the beginning of metric keys.
 
     Raises:
-        ImportError:
+        ModuleNotFoundError:
             If required TestTube package is not installed on the device.
     """
 
@@ -106,7 +106,7 @@ def __init__(
             " `pytorch_lightning.loggers.TensorBoardLogger` as an alternative."
         )
         if Experiment is None:
-            raise ImportError(
+            raise ModuleNotFoundError(
                 "You want to use `test_tube` logger which is not installed yet,"
                 " install it with `pip install test-tube`."
             )

pytorch_lightning/loggers/wandb.py

@@ -37,7 +37,7 @@
 try:
     import wandb
     from wandb.wandb_run import Run
-except ImportError:
+except ModuleNotFoundError:
     # needed for test mocks, these tests shall be updated
     wandb, Run = None, None
 
@@ -74,7 +74,7 @@ class WandbLogger(LightningLoggerBase):
         \**kwargs: Arguments passed to :func:`wandb.init` like `entity`, `group`, `tags`, etc.
 
     Raises:
-        ImportError:
+        ModuleNotFoundError:
             If required WandB package is not installed on the device.
         MisconfigurationException:
             If both ``log_model`` and ``offline``is set to ``True``.
@@ -114,7 +114,7 @@ def __init__(
         **kwargs,
     ):
         if wandb is None:
-            raise ImportError(
+            raise ModuleNotFoundError(
                 "You want to use `wandb` logger which is not installed yet,"
                 " install it with `pip install wandb`."  # pragma: no-cover
             )