Repository: microsoft/torchgeo

Issue:
Required dependencies, lazy imports
Our number of dependencies is rapidly increasing. We should think about which of these dependencies are required (`install_requires`) vs. optional (`extras_require`). Here is a proposal:

1. Required: dependencies that are needed to do just about anything with TorchGeo
2. Optional: dependencies that are only needed for optional functionality, or for a small number of datasets

This gets a bit tricky, and is currently at odds with our dependency list. For example, `rasterio` is only used in `RasterDataset`, and `fiona` is only used in `VectorDataset`. While almost half of our datasets are `RasterDataset`, ony 1 is currently `VectorDataset`. Also, `matplotlib` is only needed to plot example samples.

Here is another proposal:

1. Required: any dependency that is used in regular usage
2. Optional: dependencies used only in rare cases (single dataset)

This may be a better default. Most users will likely run `pip install torchgeo`, which will install only the things in `install_requires`. We don't want a useless installation to be the default, and `extras_require` is off by default. Alternatively, we could do a better job of documenting the recommended way to install TorchGeo and specify that you may want `pip install torchgeo[datasets,train]` or something like that.

Another thing to consider is how to handle these optional imports. We can't put them at the module level (in the case of datasets) so we use lazy imports instead. We also may want to create a wrapper like `importorraise` (akin to `pytest`'s `importorskip`) that prints a more useful error message upon `ImportError`. We could go even further and use lazy imports for almost all imports, not just optional ones. This will greatly speed up importing torchgeo.


Discussion:

Resolving patch:
diff --git a/torchgeo/datasets/idtrees.py b/torchgeo/datasets/idtrees.py
--- a/torchgeo/datasets/idtrees.py
+++ b/torchgeo/datasets/idtrees.py
@@ -522,7 +522,7 @@ def plot_las(self, index: int, colormap: Optional[str] = None) -> Any:
             import open3d  # noqa: F401
         except ImportError:
             raise ImportError(
-                "open3d is not installed and is required to use this dataset"
+                "open3d is not installed and is required to plot point clouds"
             )
         import laspy
 
diff --git a/torchgeo/datasets/nwpu.py b/torchgeo/datasets/nwpu.py
--- a/torchgeo/datasets/nwpu.py
+++ b/torchgeo/datasets/nwpu.py
@@ -117,7 +117,12 @@ def __init__(
 
         if split == "positive":
             # Must be installed to parse annotations file
-            from pycocotools.coco import COCO
+            try:
+                from pycocotools.coco import COCO  # noqa: F401
+            except ImportError:
+                raise ImportError(
+                    "pycocotools is not installed and is required to use this dataset"
+                )
 
             self.coco = COCO(
                 os.path.join(
diff --git a/torchgeo/datasets/so2sat.py b/torchgeo/datasets/so2sat.py
--- a/torchgeo/datasets/so2sat.py
+++ b/torchgeo/datasets/so2sat.py
@@ -132,7 +132,12 @@ def __init__(
             AssertionError: if ``split`` argument is invalid
             RuntimeError: if data is not found in ``root``, or checksums don't match
         """
-        import h5py
+        try:
+            import h5py  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "h5py is not installed and is required to use this dataset"
+            )
 
         assert split in ["train", "validation", "test"]
 


Test patch:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -20,9 +20,25 @@ jobs:
     - name: Install pip dependencies
       run: |
         pip install cython numpy  # needed for pycocotools
-        pip install .[datasets,tests,train]
+        pip install .[datasets,tests]
     - name: Run mypy checks
       run: mypy .
+  datasets:
+    name: datasets
+    runs-on: ubuntu-latest
+    steps:
+    - name: Clone repo
+      uses: actions/checkout@v2
+    - name: Set up python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.9
+    - name: Install pip dependencies
+      run: |
+        pip install cython numpy  # needed for pycocotools
+        pip install .[tests]
+    - name: Run pytest checks
+      run: pytest --cov=torchgeo --cov-report=xml
   pytest:
     name: pytest
     runs-on: ${{ matrix.os }}
@@ -64,7 +80,7 @@ jobs:
     - name: Install pip dependencies
       run: |
         pip install cython numpy  # needed for pycocotools
-        pip install .[datasets,tests,train]
+        pip install .[datasets,tests]
     - name: Run pytest checks
       run: pytest --cov=torchgeo --cov-report=xml
     - name: Report coverage
diff --git a/tests/datasets/test_advance.py b/tests/datasets/test_advance.py
--- a/tests/datasets/test_advance.py
+++ b/tests/datasets/test_advance.py
@@ -25,7 +25,6 @@ class TestADVANCE:
     def dataset(
         self, monkeypatch: Generator[MonkeyPatch, None, None], tmp_path: Path
     ) -> ADVANCE:
-        pytest.importorskip("scipy", minversion="0.9.0")
         monkeypatch.setattr(  # type: ignore[attr-defined]
             torchgeo.datasets.utils, "download_url", download_url
         )
@@ -57,6 +56,7 @@ def mocked_import(name: str, *args: Any, **kwargs: Any) -> Any:
         )
 
     def test_getitem(self, dataset: ADVANCE) -> None:
+        pytest.importorskip("scipy", minversion="0.9.0")
         x = dataset[0]
         assert isinstance(x, dict)
         assert isinstance(x["image"], torch.Tensor)
diff --git a/tests/datasets/test_idtrees.py b/tests/datasets/test_idtrees.py
--- a/tests/datasets/test_idtrees.py
+++ b/tests/datasets/test_idtrees.py
@@ -19,6 +19,9 @@
 import torchgeo.datasets.utils
 from torchgeo.datasets import IDTReeS
 
+pytest.importorskip("pandas", minversion="0.19.1")
+pytest.importorskip("laspy", minversion="2")
+
 
 def download_url(url: str, root: str, *args: str, **kwargs: str) -> None:
     shutil.copy(url, root)
@@ -32,8 +35,6 @@ def dataset(
         tmp_path: Path,
         request: SubRequest,
     ) -> IDTReeS:
-        pytest.importorskip("pandas")
-        pytest.importorskip("laspy")
         monkeypatch.setattr(  # type: ignore[attr-defined]
             torchgeo.datasets.idtrees, "download_url", download_url
         )
@@ -124,11 +125,11 @@ def test_mock_missing_module(
                 ImportError,
                 match=f"{package} is not installed and is required to use this dataset",
             ):
-                IDTReeS(dataset.root, download=True, checksum=True)
-        else:
+                IDTReeS(dataset.root, dataset.split, dataset.task)
+        elif package in ["open3d"]:
             with pytest.raises(
                 ImportError,
-                match=f"{package} is not installed and is required to use this dataset",
+                match=f"{package} is not installed and is required to plot point cloud",
             ):
                 dataset.plot_las(0)
 
@@ -153,7 +154,7 @@ def test_plot(self, dataset: IDTReeS) -> None:
         reason="segmentation fault on macOS and windows",
     )
     def test_plot_las(self, dataset: IDTReeS) -> None:
-        pytest.importorskip("open3d")
+        pytest.importorskip("open3d", minversion="0.11.2")
         vis = dataset.plot_las(index=0, colormap="BrBG")
         vis.close()
         vis = dataset.plot_las(index=0, colormap=None)
diff --git a/tests/datasets/test_nwpu.py b/tests/datasets/test_nwpu.py
--- a/tests/datasets/test_nwpu.py
+++ b/tests/datasets/test_nwpu.py
@@ -1,11 +1,12 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
+import builtins
 import os
 import shutil
 import sys
 from pathlib import Path
-from typing import Generator
+from typing import Any, Generator
 
 import pytest
 import torch
@@ -17,7 +18,6 @@
 import torchgeo.datasets.utils
 from torchgeo.datasets import VHR10
 
-pytest.importorskip("rarfile")
 pytest.importorskip("pycocotools")
 
 
@@ -35,6 +35,7 @@ def dataset(
         tmp_path: Path,
         request: SubRequest,
     ) -> VHR10:
+        pytest.importorskip("rarfile", minversion="3")
         monkeypatch.setattr(  # type: ignore[attr-defined]
             torchgeo.datasets.nwpu, "download_url", download_url
         )
@@ -54,6 +55,21 @@ def dataset(
         transforms = nn.Identity()  # type: ignore[attr-defined]
         return VHR10(root, split, transforms, download=True, checksum=True)
 
+    @pytest.fixture
+    def mock_missing_module(
+        self, monkeypatch: Generator[MonkeyPatch, None, None]
+    ) -> None:
+        import_orig = builtins.__import__
+
+        def mocked_import(name: str, *args: Any, **kwargs: Any) -> Any:
+            if name == "pycocotools.coco":
+                raise ImportError()
+            return import_orig(name, *args, **kwargs)
+
+        monkeypatch.setattr(  # type: ignore[attr-defined]
+            builtins, "__import__", mocked_import
+        )
+
     def test_getitem(self, dataset: VHR10) -> None:
         x = dataset[0]
         assert isinstance(x, dict)
@@ -84,3 +100,13 @@ def test_invalid_split(self) -> None:
     def test_not_downloaded(self, tmp_path: Path) -> None:
         with pytest.raises(RuntimeError, match="Dataset not found or corrupted."):
             VHR10(str(tmp_path))
+
+    def test_mock_missing_module(
+        self, dataset: VHR10, mock_missing_module: None
+    ) -> None:
+        if dataset.split == "positive":
+            with pytest.raises(
+                ImportError,
+                match="pycocotools is not installed and is required to use this datase",
+            ):
+                VHR10(dataset.root, dataset.split)
diff --git a/tests/datasets/test_resisc45.py b/tests/datasets/test_resisc45.py
--- a/tests/datasets/test_resisc45.py
+++ b/tests/datasets/test_resisc45.py
@@ -17,8 +17,6 @@
 import torchgeo.datasets.utils
 from torchgeo.datasets import RESISC45, RESISC45DataModule
 
-pytest.importorskip("rarfile")
-
 
 def download_url(url: str, root: str, *args: str, **kwargs: str) -> None:
     shutil.copy(url, root)
@@ -33,6 +31,8 @@ def dataset(
         tmp_path: Path,
         request: SubRequest,
     ) -> RESISC45:
+        pytest.importorskip("rarfile", minversion="3")
+
         monkeypatch.setattr(  # type: ignore[attr-defined]
             torchgeo.datasets.resisc45, "download_url", download_url
         )
diff --git a/tests/datasets/test_so2sat.py b/tests/datasets/test_so2sat.py
--- a/tests/datasets/test_so2sat.py
+++ b/tests/datasets/test_so2sat.py
@@ -1,9 +1,10 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
+import builtins
 import os
 from pathlib import Path
-from typing import Generator
+from typing import Any, Generator
 
 import matplotlib.pyplot as plt
 import pytest
@@ -14,6 +15,8 @@
 
 from torchgeo.datasets import So2Sat, So2SatDataModule
 
+pytest.importorskip("h5py")
+
 
 class TestSo2Sat:
     @pytest.fixture(params=["train", "validation", "test"])
@@ -32,6 +35,21 @@ def dataset(
         transforms = nn.Identity()  # type: ignore[attr-defined]
         return So2Sat(root, split, transforms, checksum=True)
 
+    @pytest.fixture
+    def mock_missing_module(
+        self, monkeypatch: Generator[MonkeyPatch, None, None]
+    ) -> None:
+        import_orig = builtins.__import__
+
+        def mocked_import(name: str, *args: Any, **kwargs: Any) -> Any:
+            if name == "h5py":
+                raise ImportError()
+            return import_orig(name, *args, **kwargs)
+
+        monkeypatch.setattr(  # type: ignore[attr-defined]
+            builtins, "__import__", mocked_import
+        )
+
     def test_getitem(self, dataset: So2Sat) -> None:
         x = dataset[0]
         assert isinstance(x, dict)
@@ -65,6 +83,15 @@ def test_plot(self, dataset: So2Sat) -> None:
         dataset.plot(x)
         plt.close()
 
+    def test_mock_missing_module(
+        self, dataset: So2Sat, mock_missing_module: None
+    ) -> None:
+        with pytest.raises(
+            ImportError,
+            match="h5py is not installed and is required to use this dataset",
+        ):
+            So2Sat(dataset.root)
+
 
 class TestSo2SatDataModule:
     @pytest.fixture(scope="class", params=zip([True, False], ["rgb", "s2"]))
diff --git a/tests/datasets/test_utils.py b/tests/datasets/test_utils.py
--- a/tests/datasets/test_utils.py
+++ b/tests/datasets/test_utils.py
@@ -98,7 +98,7 @@ def test_mock_missing_module(mock_missing_module: None) -> None:
     ],
 )
 def test_extract_archive(src: str, tmp_path: Path) -> None:
-    pytest.importorskip("rarfile")
+    pytest.importorskip("rarfile", minversion="3")
     extract_archive(os.path.join("tests", "data", src), str(tmp_path))