[App] Add utility to get install command for package extras

requirements/app/base.txt

@@ -8,6 +8,7 @@ croniter>=1.3.0, <1.4.0  # strict; TODO: for now until we find something more ro
 traitlets>=5.3.0, <=5.4.0
 arrow>=1.2.0, <1.2.4
 lightning-utilities>=0.3.*, !=0.4.0, <0.5.0
+importlib-metadata>=4.0.0; python_version < '3.8'
 beautifulsoup4>=4.8.0, <4.11.2
 inquirer>=2.10.0
 psutil<5.9.4

src/lightning_app/utilities/imports.py

@@ -17,6 +17,23 @@
 from typing import List, Union
 
 from lightning_utilities.core.imports import module_available
+from packaging.requirements import Marker, Requirement
+
+try:
+    from importlib import metadata
+except ImportError:
+    # Python < 3.8
+    import importlib_metadata as metadata  # type: ignore
+
+
+def _get_extras(extras: str = "cloud") -> List[str]:
+    """Get the list of installable packages in the given extras."""
+    from lightning_app import __package_name__
+
+    extras = f"app-{extras}" if __package_name__ == "lightning" else extras
+    requirements = {r: Requirement(r) for r in metadata.requires(__package_name__)}
+    marker = Marker(f'extra == "{extras}"')
+    return [r.split(";")[0].strip() for r, req in requirements.items() if str(req.marker) == str(marker)]
 
 
 def requires(module_paths: Union[str, List]):

tests/tests_app/utilities/test_imports.py

@@ -3,7 +3,14 @@
 
 import pytest
 
-from lightning_app.utilities.imports import requires
+from lightning_app.utilities.imports import _get_extras, requires
+
+
+@mock.patch("lightning_app.utilities.imports.metadata.requires")
+def test_get_extras(mock_requires):
+    mock_requires.return_value = ["docker (>=5.0.0) ; extra == 'test'"]
+
+    assert _get_extras("test") == ["docker (>=5.0.0)"]
 
 
 @mock.patch.dict(os.environ, {"LIGHTING_TESTING": "0"})