metaflow 2.15.20__py2.py3-none-any.whl → 2.16.0__py2.py3-none-any.whl

metaflow/packaging_sys/backend.py

@@ -0,0 +1,113 @@
+from abc import ABC, abstractmethod
+from io import BytesIO
+from typing import Any, IO, List, Optional, Union
+
+
+class PackagingBackend(ABC):
+    _mappings = {}
+    type = "none"
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if cls.type in cls._mappings:
+            raise ValueError(f"PackagingBackend {cls.type} already exists")
+        cls._mappings[cls.type] = cls
+
+    @classmethod
+    def get_backend(cls, name: str) -> "PackagingBackend":
+        if name not in cls._mappings:
+            raise ValueError(f"PackagingBackend {name} not found")
+        return cls._mappings[name]
+
+    @classmethod
+    def backend_type(cls) -> str:
+        return cls.type
+
+    @classmethod
+    @abstractmethod
+    def get_extract_commands(cls, archive_name: str, dest_dir: str) -> List[str]:
+        pass
+
+    def __init__(self):
+        self._archive = None
+
+    @abstractmethod
+    def create(self) -> "PackagingBackend":
+        pass
+
+    @abstractmethod
+    def add_file(self, filename: str, arcname: Optional[str] = None):
+        pass
+
+    @abstractmethod
+    def add_data(self, data: BytesIO, arcname: str):
+        pass
+
+    @abstractmethod
+    def close(self):
+        pass
+
+    @abstractmethod
+    def get_blob(self) -> Optional[Union[bytes, bytearray]]:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def cls_open(cls, content: IO[bytes]) -> Any:
+        """Open the archive from the given content."""
+        pass
+
+    @classmethod
+    @abstractmethod
+    def cls_has_member(cls, archive: Any, name: str) -> bool:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def cls_get_member(cls, archive: Any, name: str) -> Optional[bytes]:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def cls_extract_members(
+        cls,
+        archive: Any,
+        members: Optional[List[str]] = None,
+        dest_dir: str = ".",
+    ) -> None:
+        pass
+
+    @classmethod
+    @abstractmethod
+    def cls_list_members(cls, archive: Any) -> Optional[List[str]]:
+        pass
+
+    def has_member(self, name: str) -> bool:
+        if self._archive:
+            return self.cls_has_member(self._archive, name)
+        raise ValueError("Cannot check for member in an uncreated archive")
+
+    def get_member(self, name: str) -> Optional[bytes]:
+        if self._archive:
+            return self.cls_get_member(self._archive, name)
+        raise ValueError("Cannot get member from an uncreated archive")
+
+    def extract_members(
+        self, members: Optional[List[str]] = None, dest_dir: str = "."
+    ) -> None:
+        if self._archive:
+            self.cls_extract_members(self._archive, members, dest_dir)
+        else:
+            raise ValueError("Cannot extract from an uncreated archive")
+
+    def list_members(self) -> Optional[List[str]]:
+        if self._archive:
+            return self.cls_list_members(self._archive)
+        raise ValueError("Cannot list members from an uncreated archive")
+
+    def __enter__(self):
+        self.create()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()

metaflow/packaging_sys/distribution_support.py

@@ -0,0 +1,153 @@
+# Support saving of distribution information so we can give it back to users even
+# if we do not install those distributions. This is used to package distributions in
+# the MetaflowCodeContent package and provide an experience as if the packages were installed
+# system-wide.
+
+import os
+import re
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import (
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    NamedTuple,
+    Optional,
+    Set,
+    TYPE_CHECKING,
+    Union,
+    cast,
+)
+
+import inspect
+from collections import defaultdict
+
+from ..extension_support import metadata
+from ..util import get_metaflow_root
+
+if TYPE_CHECKING:
+    import pathlib
+
+_cached_distributions = None
+
+packages_distributions = None  # type: Optional[Callable[[], Mapping[str, List[str]]]]
+name_normalizer = re.compile(r"[-_.]+")
+
+if sys.version_info[:2] >= (3, 10):
+    packages_distributions = metadata.packages_distributions
+else:
+    # This is the code present in 3.10+ -- we replicate here for other versions
+    def _packages_distributions() -> Mapping[str, List[str]]:
+        """
+        Return a mapping of top-level packages to their
+        distributions.
+        """
+        pkg_to_dist = defaultdict(list)
+        for dist in metadata.distributions():
+            for pkg in _top_level_declared(dist) or _top_level_inferred(dist):
+                pkg_to_dist[pkg].append(dist.metadata["Name"])
+        return dict(pkg_to_dist)
+
+    def _top_level_declared(dist: metadata.Distribution) -> List[str]:
+        return (dist.read_text("top_level.txt") or "").split()
+
+    def _topmost(name: "pathlib.PurePosixPath") -> Optional[str]:
+        """
+        Return the top-most parent as long as there is a parent.
+        """
+        top, *rest = name.parts
+        return top if rest else None
+
+    def _get_toplevel_name(name: "pathlib.PurePosixPath") -> str:
+        return _topmost(name) or (
+            # python/typeshed#10328
+            inspect.getmodulename(name)  # type: ignore
+            or str(name)
+        )
+
+    def _top_level_inferred(dist: "metadata.Distribution"):
+        opt_names = set(map(_get_toplevel_name, dist.files or []))
+
+        def importable_name(name):
+            return "." not in name
+
+        return filter(importable_name, opt_names)
+
+    packages_distributions = _packages_distributions
+
+
+def modules_to_distributions() -> Dict[str, List[metadata.Distribution]]:
+    """
+    Return a mapping of top-level modules to their distributions.
+
+    Returns
+    -------
+    Dict[str, List[metadata.Distribution]]
+        A mapping of top-level modules to their distributions.
+    """
+    global _cached_distributions
+    pd = cast(Callable[[], Mapping[str, List[str]]], packages_distributions)
+    if _cached_distributions is None:
+        _cached_distributions = {
+            k: [metadata.distribution(d) for d in v] for k, v in pd().items()
+        }
+    return _cached_distributions
+
+
+_ModuleInfo = NamedTuple(
+    "_ModuleInfo",
+    [
+        ("name", str),
+        ("root_paths", Set[str]),
+        ("module", ModuleType),
+        ("metaflow_module", bool),
+    ],
+)
+
+
+class PackagedDistribution(metadata.Distribution):
+    """
+    A Python Package packaged within a MetaflowCodeContent. This allows users to use use importlib
+    as they would regularly and the packaged Python Package would be considered as a
+    distribution even if it really isn't (since it is just included in the PythonPath).
+    """
+
+    def __init__(self, root: str, content: Dict[str, str]):
+        self._root = Path(root)
+        self._content = content
+
+    # Strongly inspired from PathDistribution in metadata.py
+    def read_text(self, filename: Union[str, os.PathLike]) -> Optional[str]:
+        if str(filename) in self._content:
+            return self._content[str(filename)]
+        return None
+
+    read_text.__doc__ = metadata.Distribution.read_text.__doc__
+
+    # Returns a metadata.SimplePath but not always present in importlib.metadata libs so
+    # skipping return type.
+    def locate_file(self, path: Union[str, os.PathLike]):
+        return self._root / path
+
+
+class PackagedDistributionFinder(metadata.DistributionFinder):
+    def __init__(self, dist_info: Dict[str, Dict[str, str]]):
+        self._dist_info = dist_info
+
+    def find_distributions(self, context=metadata.DistributionFinder.Context()):
+        if context.name is None:
+            # Yields all known distributions
+            for name, info in self._dist_info.items():
+                yield PackagedDistribution(
+                    os.path.join(get_metaflow_root(), name), info
+                )
+            return None
+        name = name_normalizer.sub("-", cast(str, context.name)).lower()
+        if name in self._dist_info:
+            yield PackagedDistribution(
+                os.path.join(get_metaflow_root(), cast(str, context.name)),
+                self._dist_info[name],
+            )
+        return None

metaflow/packaging_sys/tar_backend.py

@@ -0,0 +1,86 @@
+import tarfile
+
+from io import BytesIO
+from typing import IO, List, Optional, Union
+
+from .backend import PackagingBackend
+
+
+class TarPackagingBackend(PackagingBackend):
+    type = "tgz"
+
+    @classmethod
+    def get_extract_commands(cls, archive_name: str, dest_dir: str) -> List[str]:
+        return [
+            f"TAR_OPTIONS='--warning=no-timestamp' tar -xzf {archive_name} -C {dest_dir}"
+        ]
+
+    def __init__(self):
+        super().__init__()
+        self._buf = None
+
+    def create(self):
+        self._buf = BytesIO()
+        self._archive = tarfile.open(
+            fileobj=self._buf, mode="w:gz", compresslevel=3, dereference=True
+        )
+        return self
+
+    def add_file(self, filename: str, arcname: Optional[str] = None):
+        info = self._archive.gettarinfo(filename, arcname)
+        # Setting this default to Dec 3, 2019
+        info.mtime = 1575360000
+        with open(filename, mode="rb") as f:
+            self._archive.addfile(info, f)
+
+    def add_data(self, data: BytesIO, arcname: str):
+        info = tarfile.TarInfo(arcname)
+        data.seek(0)
+        info.size = len(data.getvalue())
+        # Setting this default to Dec 3, 2019
+        info.mtime = 1575360000
+        self._archive.addfile(info, data)
+
+    def close(self):
+        if self._archive:
+            self._archive.close()
+
+    def get_blob(self) -> Optional[Union[bytes, bytearray]]:
+        if self._buf:
+            blob = bytearray(self._buf.getvalue())
+            blob[4:8] = [0] * 4  # Reset 4 bytes from offset 4 to account for ts
+            return blob
+        return None
+
+    @classmethod
+    def cls_open(cls, content: IO[bytes]) -> tarfile.TarFile:
+        return tarfile.open(fileobj=content, mode="r:gz")
+
+    @classmethod
+    def cls_has_member(cls, archive: tarfile.TarFile, name: str) -> bool:
+        try:
+            archive.getmember(name)
+            return True
+        except KeyError:
+            return False
+
+    @classmethod
+    def cls_get_member(cls, archive: tarfile.TarFile, name: str) -> Optional[bytes]:
+        try:
+            member = archive.getmember(name)
+            return archive.extractfile(member).read()
+        except KeyError:
+            return None
+
+    @classmethod
+    def cls_extract_members(
+        cls,
+        archive: tarfile.TarFile,
+        members: Optional[List[str]] = None,
+        dest_dir: str = ".",
+    ) -> None:
+        archive.extractall(path=dest_dir, members=members)
+
+    @classmethod
+    def cls_list_members(cls, archive: tarfile.TarFile) -> Optional[List[str]]:
+        return archive.getnames() or None

metaflow/packaging_sys/utils.py

@@ -0,0 +1,91 @@
+import os
+from contextlib import contextmanager
+from typing import Callable, Generator, List, Optional, Tuple
+
+from ..util import to_unicode
+
+
+# this is os.walk(follow_symlinks=True) with cycle detection
+def walk_without_cycles(
+    top_root: str,
+    exclude_dirs: Optional[List[str]] = None,
+) -> Generator[Tuple[str, List[str]], None, None]:
+    seen = set()
+
+    default_skip_dirs = ["__pycache__"]
+
+    def _recurse(root, skip_dirs):
+        for parent, dirs, files in os.walk(root):
+            dirs[:] = [d for d in dirs if d not in skip_dirs]
+            for d in dirs:
+                path = os.path.join(parent, d)
+                if os.path.islink(path):
+                    # Breaking loops: never follow the same symlink twice
+                    #
+                    # NOTE: this also means that links to sibling links are
+                    # not followed. In this case:
+                    #
+                    #   x -> y
+                    #   y -> oo
+                    #   oo/real_file
+                    #
+                    # real_file is only included twice, not three times
+                    reallink = os.path.realpath(path)
+                    if reallink not in seen:
+                        seen.add(reallink)
+                        for x in _recurse(path, default_skip_dirs):
+                            yield x
+            yield parent, files
+
+    skip_dirs = set(default_skip_dirs + (exclude_dirs or []))
+    for x in _recurse(top_root, skip_dirs):
+        skip_dirs = default_skip_dirs
+        yield x
+
+
+def walk(
+    root: str,
+    exclude_hidden: bool = True,
+    file_filter: Optional[Callable[[str], bool]] = None,
+    exclude_tl_dirs: Optional[List[str]] = None,
+) -> Generator[Tuple[str, str], None, None]:
+    root = to_unicode(root)  # handle files/folder with non ascii chars
+    prefixlen = len("%s/" % os.path.dirname(root))
+    for (
+        path,
+        files,
+    ) in walk_without_cycles(root, exclude_tl_dirs):
+        if exclude_hidden and "/." in path:
+            continue
+        # path = path[2:] # strip the ./ prefix
+        # if path and (path[0] == '.' or './' in path):
+        #    continue
+        for fname in files:
+            if file_filter is None or file_filter(fname):
+                p = os.path.join(path, fname)
+                yield p, p[prefixlen:]
+
+
+def suffix_filter(suffixes: List[str]) -> Callable[[str], bool]:
+    """
+    Returns a filter function that checks if a file ends with any of the given suffixes.
+    """
+    suffixes = [s.lower() for s in suffixes]
+
+    def _filter(fname: str) -> bool:
+        fname = fname.lower()
+        return (
+            suffixes is None
+            or (fname[0] == "." and fname in suffixes)
+            or (fname[0] != "." and any(fname.endswith(suffix) for suffix in suffixes))
+        )
+
+    return _filter
+
+
+@contextmanager
+def with_dir(new_dir):
+    current_dir = os.getcwd()
+    os.chdir(new_dir)
+    yield new_dir
+    os.chdir(current_dir)