Skip to content

Commit a881db9

Browse files
ethanwharrisBorda
authored andcommitted
[App] Add utility to get install command for package extras (#15809)
(cherry picked from commit f171657)
1 parent f844e19 commit a881db9

File tree

2 files changed

+35
-1
lines changed

2 files changed

+35
-1
lines changed

src/lightning_app/utilities/imports.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,30 @@
1717
from typing import List, Union
1818

1919
from lightning_utilities.core.imports import module_available
20+
from packaging.requirements import Marker, Requirement
21+
22+
try:
23+
from importlib import metadata
24+
except ImportError:
25+
# Python < 3.8
26+
import importlib_metadata as metadata # type: ignore
27+
28+
29+
def _get_extras(extras: str) -> str:
30+
"""Get the given extras as a space delimited string.
31+
32+
Used by the platform to install cloud extras in the cloud.
33+
"""
34+
from lightning_app import __package_name__
35+
36+
requirements = {r: Requirement(r) for r in metadata.requires(__package_name__)}
37+
marker = Marker(f'extra == "{extras}"')
38+
requirements = [r for r, req in requirements.items() if str(req.marker) == str(marker)]
39+
40+
if requirements:
41+
requirements = [f"'{r.split(';')[0].strip()}'" for r in requirements]
42+
return " ".join(requirements)
43+
return ""
2044

2145

2246
def requires(module_paths: Union[str, List]):

tests/tests_app/utilities/test_imports.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,17 @@
33

44
import pytest
55

6-
from lightning_app.utilities.imports import requires
6+
from lightning_app import __package_name__
7+
from lightning_app.utilities.imports import _get_extras, requires
8+
9+
10+
def test_get_extras():
11+
extras = "app-cloud" if __package_name__ == "lightning" else "cloud"
12+
extras = _get_extras(extras)
13+
assert "docker" in extras
14+
assert "redis" in extras
15+
16+
assert _get_extras("fake-extras") == ""
717

818

919
@mock.patch.dict(os.environ, {"LIGHTING_TESTING": "0"})

0 commit comments

Comments
 (0)