Skip to content

Commit

Permalink
Merge pull request #33487 Better error message on incorrect package l…
Browse files Browse the repository at this point in the history
…isting.
  • Loading branch information
robertwb authored Jan 3, 2025
2 parents f27547d + 4c14921 commit 98c8534
Showing 1 changed file with 16 additions and 9 deletions.
25 changes: 16 additions & 9 deletions sdks/python/apache_beam/yaml/yaml_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
from typing import Dict
from typing import Iterable
from typing import Iterator
from typing import List
from typing import Mapping
from typing import Optional

Expand Down Expand Up @@ -349,7 +350,7 @@ def python(urns, packages=()):

@ExternalProvider.register_provider_type('pythonPackage')
class ExternalPythonProvider(ExternalProvider):
def __init__(self, urns, packages):
def __init__(self, urns, packages: Iterable[str]):
super().__init__(urns, PypiExpansionService(packages))

def available(self):
Expand Down Expand Up @@ -1017,26 +1018,31 @@ class PypiExpansionService:
"""
VENV_CACHE = os.path.expanduser("~/.apache_beam/cache/venvs")

def __init__(self, packages, base_python=sys.executable):
self._packages = packages
def __init__(
self, packages: Iterable[str], base_python: str = sys.executable):
if not isinstance(packages, Iterable) or isinstance(packages, str):
raise TypeError(
"Packages must be an iterable of strings, got %r" % packages)
self._packages = list(packages)
self._base_python = base_python

@classmethod
def _key(cls, base_python, packages):
def _key(cls, base_python: str, packages: List[str]) -> str:
return json.dumps({
'binary': base_python, 'packages': sorted(packages)
},
sort_keys=True)

@classmethod
def _path(cls, base_python, packages):
def _path(cls, base_python: str, packages: List[str]) -> str:
return os.path.join(
cls.VENV_CACHE,
hashlib.sha256(cls._key(base_python,
packages).encode('utf-8')).hexdigest())

@classmethod
def _create_venv_from_scratch(cls, base_python, packages):
def _create_venv_from_scratch(
cls, base_python: str, packages: List[str]) -> str:
venv = cls._path(base_python, packages)
if not os.path.exists(venv):
try:
Expand All @@ -1054,7 +1060,8 @@ def _create_venv_from_scratch(cls, base_python, packages):
return venv

@classmethod
def _create_venv_from_clone(cls, base_python, packages):
def _create_venv_from_clone(
cls, base_python: str, packages: List[str]) -> str:
venv = cls._path(base_python, packages)
if not os.path.exists(venv):
try:
Expand All @@ -1074,7 +1081,7 @@ def _create_venv_from_clone(cls, base_python, packages):
return venv

@classmethod
def _create_venv_to_clone(cls, base_python):
def _create_venv_to_clone(cls, base_python: str) -> str:
if '.dev' in beam_version:
base_venv = os.path.dirname(os.path.dirname(base_python))
print('Cloning dev environment from', base_venv)
Expand All @@ -1085,7 +1092,7 @@ def _create_venv_to_clone(cls, base_python):
'virtualenv-clone'
])

def _venv(self):
def _venv(self) -> str:
return self._create_venv_from_clone(self._base_python, self._packages)

def __enter__(self):
Expand Down

0 comments on commit 98c8534

Please sign in to comment.