add optional requirements to setuptools installation

.circleci/config.yml.in

@@ -144,13 +144,24 @@ commands:
       editable:
         type: boolean
         default: true
+      optional_dependencies:
+        type: boolean
+        default: true
     steps:
       - pip_install:
-          args: --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
           descr: Install PyTorch from nightly releases
+          args: --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+
       - pip_install:
-          args: --no-build-isolation <<# parameters.editable >> --editable <</ parameters.editable >> .
-          descr: Install torchvision <<# parameters.editable >> in editable mode <</ parameters.editable >>
+          descr: >
+            Install torchvision
+            <<# parameters.editable >> in editable mode <</ parameters.editable >>
+            <<# parameters.optional_dependencies >> with optional dependencies <</ parameters.optional_dependencies >>
+          args: >
+            --no-build-isolation
+            <<# parameters.editable >> --editable <</ parameters.editable >>
+            .<<# parameters.optional_dependencies >>[all]<</ parameters.optional_dependencies >>
+
 
   install_prototype_dependencies:
     steps:
@@ -350,9 +361,6 @@ jobs:
                 | parallel -j0 'wget --no-verbose -O ~/.cache/torch/hub/checkpoints/`basename {}` {}\?source=ci'
       - install_torchvision
       - install_prototype_dependencies
-      - pip_install:
-          args: scipy pycocotools
-          descr: Install optional dependencies
       - run:
           name: Enable prototype tests
           command: echo 'export PYTORCH_TEST_WITH_PROTOTYPE=1' >> $BASH_ENV

.circleci/unittest/linux/scripts/environment.yml

@@ -11,6 +11,3 @@ dependencies:
   - ca-certificates
   - pip:
     - future
-    - pillow >=5.3.0, !=8.3.*
-    - scipy
-    - av

.circleci/unittest/linux/scripts/install.sh

@@ -42,5 +42,8 @@ if [ $PYTHON_VERSION == "3.6" ]; then
     pip install "pillow>=5.3.0,!=8.3.*"
 fi
 
+printf "* Installing optional dependencies\n"
+pip install -r optional-requirements.txt
+
 printf "* Installing torchvision\n"
 python setup.py develop

.circleci/unittest/windows/scripts/environment.yml

@@ -11,7 +11,4 @@ dependencies:
   - ca-certificates
   - pip:
     - future
-    - pillow >=5.3.0, !=8.3.*
-    - scipy
-    - av
     - dataclasses

.circleci/unittest/windows/scripts/install.sh

@@ -48,5 +48,8 @@ fi
 
 source "$this_dir/set_cuda_envs.sh"
 
+printf "* Installing optional dependencies\n"
+pip install -r optional-requirements.txt
+
 printf "* Installing torchvision\n"
 "$this_dir/vc_env_helper.bat" python setup.py develop

optional-requirements.txt

@@ -0,0 +1,8 @@
+# scipy
+scipy
+# io
+av
+# datasets
+scipy
+lmdb
+pycocotools

setup.py

@@ -5,6 +5,7 @@
 import shutil
 import subprocess
 import sys
+from collections import defaultdict
 
 import torch
 from pkg_resources import parse_version, get_distribution, DistributionNotFound
@@ -444,6 +445,31 @@ def run(self):
         distutils.command.clean.clean.run(self)
 
 
+def get_optional_requirements(file_name="optional-requirements.txt"):
+    extras_require = defaultdict(lambda: set())
+    with open(os.path.join(cwd, file_name)) as file:
+        extra = None
+        for line in file:
+            if line.startswith("#"):
+                extra = line[1:].strip()
+                continue
+
+            extras_require[extra].add(line.strip())
+
+    unknown_extra = extras_require.pop(None, None)
+    if unknown_extra:
+        raise RuntimeError(
+            f"Requirement(s) {sorted(unknown_extra)} do not belong to any extra. "
+            f"Please place them in a existing extra or create new one by adding a comment with the name of the extra."
+        )
+
+    extras_require = {extra: sorted(requirements) for extra, requirements in extras_require.items()}
+    extras_require["all"] = sorted(
+        {requirement for requirements in extras_require.values() for requirement in requirements}
+    )
+    return extras_require
+
+
 if __name__ == "__main__":
     print(f"Building wheel {package_name}-{version}")
 
@@ -467,9 +493,7 @@ def run(self):
         package_data={package_name: ["*.dll", "*.dylib", "*.so", "prototype/datasets/_builtin/*.categories"]},
         zip_safe=False,
         install_requires=requirements,
-        extras_require={
-            "scipy": ["scipy"],
-        },
+        extras_require=get_optional_requirements(),
         ext_modules=get_extensions(),
         python_requires=">=3.6",
         cmdclass={

test/test_prototype_builtin_datasets.py

@@ -1,4 +1,5 @@
 import io
+import pathlib
 
 import builtin_dataset_mocks
 import pytest
@@ -51,6 +52,34 @@ def dataset_parametrization(*names, decoder=to_bytes):
     )
 
 
+def test_optional_dependencies():
+    required_dependencies = set()
+    for name in datasets.list():
+        required_dependencies.update(datasets.info(name).dependecies)
+
+    with open(pathlib.Path(__file__).parents[1] / "optional-requirements.txt") as file:
+        for line in file:
+            if line.startswith("#") and line[1:].strip() == "datasets":
+                break
+
+        registered_dependencies = set()
+        for line in file:
+            if line.startswith("#"):
+                break
+
+            registered_dependencies.add(line.strip())
+
+    # TODO: Until we have feature parity between the stable and prototype datasets, we check if the registered
+    #  dependencies are a superset of the requires ones. Afterwards we should check for equality.
+    missing_dependencies = required_dependencies - registered_dependencies
+    if missing_dependencies:
+        raise AssertionError(
+            f"The datasets depend on the third-party packages {sequence_to_str(sorted(missing_dependencies))}, "
+            f"but they are not listed in the '# datasets' section in the file $TORCHVISION/optional-requirements.txt. "
+            "Please add them."
+        )
+
+
 class TestCommon:
     @dataset_parametrization()
     def test_smoke(self, dataset, mock_info):