Skip to content

Commit

Permalink
[App] Add utility to get install command for package extras (#15809)
Browse files Browse the repository at this point in the history
  • Loading branch information
ethanwharris authored Nov 24, 2022
1 parent 8ee889b commit f171657
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 1 deletion.
24 changes: 24 additions & 0 deletions src/lightning_app/utilities/imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,30 @@
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) -> str:
"""Get the given extras as a space delimited string.
Used by the platform to install cloud extras in the cloud.
"""
from lightning_app import __package_name__

requirements = {r: Requirement(r) for r in metadata.requires(__package_name__)}
marker = Marker(f'extra == "{extras}"')
requirements = [r for r, req in requirements.items() if str(req.marker) == str(marker)]

if requirements:
requirements = [f"'{r.split(';')[0].strip()}'" for r in requirements]
return " ".join(requirements)
return ""


def requires(module_paths: Union[str, List]):
Expand Down
12 changes: 11 additions & 1 deletion tests/tests_app/utilities/test_imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,17 @@

import pytest

from lightning_app.utilities.imports import requires
from lightning_app import __package_name__
from lightning_app.utilities.imports import _get_extras, requires


def test_get_extras():
extras = "app-cloud" if __package_name__ == "lightning" else "cloud"
extras = _get_extras(extras)
assert "docker" in extras
assert "redis" in extras

assert _get_extras("fake-extras") == ""


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

0 comments on commit f171657

Please sign in to comment.