Skip to content

Commit

Permalink
Reuse assistant code to create the mirror package (#15365)
Browse files Browse the repository at this point in the history
  • Loading branch information
@carmocca
carmocca committed Nov 8, 2022
1 parent 35b66fd commit e33d09a
Show file tree
Hide file tree
Showing 6 changed files with 80 additions and 243 deletions.
217 changes: 70 additions & 147 deletions .actions/assistant.py
View file
Original file line number Diff line number Diff line change
@@ -1,22 +1,11 @@
import datetime
import glob
import json
import os
import re
import shutil
from distutils.version import LooseVersion
from importlib.util import module_from_spec, spec_from_file_location
from itertools import chain
from pathlib import Path
from pprint import pprint
from types import ModuleType
from typing import List, Optional, Sequence
from urllib import request
from urllib.request import Request, urlopen
from typing import Dict, List, Optional, Sequence, Tuple

import jsonargparse
import pkg_resources
from packaging.version import parse as version_parse

REQUIREMENT_FILES = {
"pytorch": (
Expand All @@ -36,31 +25,6 @@
),
}
REQUIREMENT_FILES_ALL = list(chain(*REQUIREMENT_FILES.values()))
PACKAGE_MAPPING = {"app": "lightning-app", "pytorch": "pytorch-lightning"}


def pypi_versions(package_name: str, drop_pre: bool = True) -> List[str]:
"""Return a list of released versions of a provided pypi name.
>>> _ = pypi_versions("lightning_app", drop_pre=False)
"""
# https://stackoverflow.com/a/27239645/4521646
url = f"https://pypi.org/pypi/{package_name}/json"
data = json.load(urlopen(Request(url)))
versions = list(data["releases"].keys())
# todo: drop this line after cleaning Pypi history from invalid versions
versions = list(filter(lambda v: v.count(".") == 2, versions))
if drop_pre:
versions = list(filter(lambda v: all(c not in v for c in ["rc", "dev"]), versions))
versions.sort(key=version_parse)
return versions


def _load_py_module(name: str, location: str) -> ModuleType:
spec = spec_from_file_location(name, location)
py = module_from_spec(spec)
spec.loader.exec_module(py)
return py


def _retrieve_files(directory: str, *ext: str) -> List[str]:
Expand All @@ -73,25 +37,74 @@ def _retrieve_files(directory: str, *ext: str) -> List[str]:
return all_files


class AssistantCLI:
_PATH_ROOT = str(Path(__file__).parent.parent)
_PATH_SRC = os.path.join(_PATH_ROOT, "src")
def _replace_imports(lines: List[str], mapping: List[Tuple[str, str]]) -> List[str]:
"""Replace imports of standalone package to lightning.
>>> lns = [
... "lightning_app",
... "delete_cloud_lightning_apps",
... "from lightning_app import",
... "lightning_apps = []",
... "lightning_app and pytorch_lightning are ours",
... "def _lightning_app():",
... ":class:`~lightning_app.core.flow.LightningFlow`"
... ]
>>> mapping = [("lightning_app", "lightning.app"), ("pytorch_lightning", "lightning.pytorch")]
>>> _replace_imports(lns, mapping) # doctest: +NORMALIZE_WHITESPACE
['lightning.app', 'delete_cloud_lightning_apps', 'from lightning.app import', 'lightning_apps = []',\
'lightning.app and lightning.pytorch are ours', 'def _lightning_app():',\
':class:`~lightning.app.core.flow.LightningFlow`']
"""
out = lines[:]
for source_import, target_import in mapping:
for i, ln in enumerate(out):
out[i] = re.sub(rf"([^_]|^){source_import}([^_\w]|$)", rf"\1{target_import}\2", ln)
return out


def copy_replace_imports(
source_dir: str, source_imports: List[str], target_imports: List[str], target_dir: Optional[str] = None
) -> None:
print(f"Replacing imports: {locals()}")
assert len(source_imports) == len(target_imports), (
"source and target imports must have the same length, "
f"source: {len(source_imports)}, target: {len(target_imports)}"
)
if target_dir is None:
target_dir = source_dir

ls = _retrieve_files(source_dir)
for fp in ls:
if fp.endswith(".py") or not fp.endswith(".pyc"):
with open(fp, encoding="utf-8") as fo:
try:
lines = fo.readlines()
except UnicodeDecodeError:
# a binary file, skip
print(f"Skipped replacing imports for {fp}")
continue
lines = _replace_imports(lines, list(zip(source_imports, target_imports)))
fp_new = fp.replace(source_dir, target_dir)
os.makedirs(os.path.dirname(fp_new), exist_ok=True)
with open(fp_new, "w", encoding="utf-8") as fo:
fo.writelines(lines)


def create_mirror_package(source_dir: str, package_mapping: Dict[str, str]) -> None:
# replace imports and copy the code
mapping = package_mapping.copy()
mapping.pop("lightning", None) # pop this key to avoid replacing `lightning` to `lightning.lightning`
for new, previous in mapping.items():
copy_replace_imports(
source_dir=os.path.join(source_dir, previous),
# pytorch_lightning uses lightning_lite, so we need to replace all imports for all directories
source_imports=list(mapping.values()),
target_imports=[f"lightning.{new}" for new in mapping],
target_dir=os.path.join(source_dir, "lightning", new),
)

@staticmethod
def prepare_nightly_version(proj_root: str = _PATH_ROOT) -> None:
"""Replace semantic version by date."""
path_info = os.path.join(proj_root, "pytorch_lightning", "__about__.py")
# get today date
now = datetime.datetime.now()
now_date = now.strftime("%Y%m%d")

print(f"prepare init '{path_info}' - replace version by {now_date}")
with open(path_info, encoding="utf-8") as fp:
init = fp.read()
init = re.sub(r'__version__ = [\d\.\w\'"]+', f'__version__ = "{now_date}"', init)
with open(path_info, "w", encoding="utf-8") as fp:
fp.write(init)

class AssistantCLI:
@staticmethod
def requirements_prune_pkgs(packages: Sequence[str], req_files: Sequence[str] = REQUIREMENT_FILES_ALL) -> None:
"""Remove some packages from given requirement files."""
Expand Down Expand Up @@ -130,107 +143,17 @@ def replace_oldest_ver(requirement_fnames: Sequence[str] = REQUIREMENT_FILES_ALL
for fname in requirement_fnames:
AssistantCLI._replace_min(fname)

@staticmethod
def _release_pkg(pkg: str, src_folder: str = _PATH_SRC) -> bool:
pypi_ver = pypi_versions(pkg)[-1]
_version = _load_py_module("version", os.path.join(src_folder, pkg.replace("-", "_"), "__version__.py"))
local_ver = _version.version
return "dev" not in local_ver and LooseVersion(local_ver) > LooseVersion(pypi_ver)

@staticmethod
def determine_releasing_pkgs(
src_folder: str = _PATH_SRC, packages: Sequence[str] = ("pytorch", "app"), inverse: bool = False
) -> Sequence[str]:
"""Determine version of package where the name is `lightning.<name>`."""
if isinstance(packages, str):
packages = [packages]
releasing = [pkg for pkg in packages if AssistantCLI._release_pkg(PACKAGE_MAPPING[pkg], src_folder=src_folder)]
if inverse:
releasing = list(filter(lambda pkg: pkg not in releasing, packages))
return json.dumps([{"pkg": pkg for pkg in releasing}])

@staticmethod
def download_package(package: str, folder: str = ".", version: Optional[str] = None) -> None:
"""Download specific or latest package from PyPI where the name is `lightning.<name>`."""
url = f"https://pypi.org/pypi/{PACKAGE_MAPPING[package]}/json"
data = json.load(urlopen(Request(url)))
if not version:
pypi_vers = pypi_versions(PACKAGE_MAPPING[package], drop_pre=False)
version = pypi_vers[-1]
releases = list(filter(lambda r: r["packagetype"] == "sdist", data["releases"][version]))
assert releases, f"Missing 'sdist' for this package/version aka {package}/{version}"
release = releases[0]
pkg_url = release["url"]
pkg_file = os.path.basename(pkg_url)
pkg_path = os.path.join(folder, pkg_file)
os.makedirs(folder, exist_ok=True)
print(f"downloading: {pkg_url}")
request.urlretrieve(pkg_url, pkg_path)

@staticmethod
def _find_pkgs(folder: str, pkg_pattern: str = "lightning") -> List[str]:
"""Find all python packages with spec.
pattern in given folder, in case `src` exists dive there.
"""
pkg_dirs = [d for d in glob.glob(os.path.join(folder, "*")) if os.path.isdir(d)]
if "src" in [os.path.basename(p) for p in pkg_dirs]:
return AssistantCLI._find_pkgs(os.path.join(folder, "src"), pkg_pattern)
pkg_dirs = list(filter(lambda p: pkg_pattern in os.path.basename(p), pkg_dirs))
return pkg_dirs

@staticmethod
def mirror_pkg2source(pypi_folder: str, src_folder: str) -> None:
"""From extracted sdist packages overwrite the python package with given pkg pattern."""
pypi_dirs = [d for d in glob.glob(os.path.join(pypi_folder, "*")) if os.path.isdir(d)]
for pkg_dir in pypi_dirs:
for py_dir in AssistantCLI._find_pkgs(pkg_dir):
dir_name = os.path.basename(py_dir)
py_dir2 = os.path.join(src_folder, dir_name)
shutil.rmtree(py_dir2, ignore_errors=True)
shutil.copytree(py_dir, py_dir2)

@staticmethod
def copy_replace_imports(
source_dir: str, source_import: str, target_import: str, target_dir: Optional[str] = None
) -> None:
"""Recursively replace imports in given folder."""

source_imports = source_import.strip().split(",")
target_imports = target_import.strip().split(",")
assert len(source_imports) == len(target_imports), (
"source and target imports must have the same length, "
f"source: {len(source_import)}, target: {len(target_import)}"
)

if target_dir is None:
target_dir = source_dir

ls = _retrieve_files(source_dir)

for fp in ls:
if fp.endswith(".py"):
with open(fp, encoding="utf-8") as fo:
py = fo.readlines()

for source_import, target_import in zip(source_imports, target_imports):
for i, ln in enumerate(py):
py[i] = re.sub(rf"([^_]|^){source_import}([^_\w]|$)", rf"\1{target_import}\2", ln)

if target_dir:
fp_new = fp.replace(source_dir, target_dir)
os.makedirs(os.path.dirname(fp_new), exist_ok=True)
else:
fp_new = fp

with open(fp_new, "w", encoding="utf-8") as fo:
fo.writelines(py)
elif not fp.endswith(".pyc"):
fp_new = fp.replace(source_dir, target_dir)
os.makedirs(os.path.dirname(fp_new), exist_ok=True)
if os.path.abspath(fp) != os.path.abspath(fp_new):
shutil.copy2(fp, fp_new)
copy_replace_imports(source_dir, source_imports, target_imports, target_dir=target_dir)


if __name__ == "__main__":
import jsonargparse

jsonargparse.CLI(AssistantCLI, as_positional=False)
4 changes: 1 addition & 3 deletions .actions/requirements.txt
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,2 @@
jsonargparse>=4.15.0
packaging
jsonargparse>=4.16.0
requests
typing_extensions
89 changes: 1 addition & 88 deletions .actions/setup_tools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,29 +20,17 @@
import tempfile
import urllib.request
from distutils.version import LooseVersion
from importlib.util import module_from_spec, spec_from_file_location
from itertools import chain
from types import ModuleType
from typing import Dict, List
from typing import List

from pkg_resources import parse_requirements

_PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))
_PACKAGE_MAPPING = {"pytorch": "pytorch_lightning", "app": "lightning_app", "lite": "lightning_lite"}

# TODO: remove this once lightning-ui package is ready as a dependency
_LIGHTNING_FRONTEND_RELEASE_URL = "https://storage.googleapis.com/grid-packages/lightning-ui/v0.0.0/build.tar.gz"


def _load_py_module(name: str, location: str) -> ModuleType:
spec = spec_from_file_location(name, location)
assert spec, f"Failed to load module {name} from {location}"
py = module_from_spec(spec)
assert spec.loader, f"ModuleSpec.loader is None for {name} from {location}"
spec.loader.exec_module(py)
return py


def _augment_requirement(ln: str, comment_char: str = "#", unfreeze: str = "all") -> str:
"""Adjust the upper version contrains.
Expand Down Expand Up @@ -163,81 +151,6 @@ def load_readme_description(path_dir: str, homepage: str, version: str) -> str:
return text


def parse_version_from_file(pkg_root: str) -> str:
"""Loading the package version from file."""
file_ver = os.path.join(pkg_root, "__version__.py")
file_about = os.path.join(pkg_root, "__about__.py")
if os.path.isfile(file_ver):
ver = _load_py_module("version", file_ver).version
elif os.path.isfile(file_about):
ver = _load_py_module("about", file_about).__version__
else: # this covers case you have build only meta-package so not additional source files are present
ver = ""
return ver


def _replace_imports_in_file(lines: List[str], pkg_lut: Dict[str, str]) -> List[str]:
"""Replace imports of standalone package to lightning.
>>> lns = ["lightning_app",
... "delete_cloud_lightning_apps",
... "from lightning_app import",
... "lightning_apps = []",
... "lightning_app is ours",
... "def _lightning_app():",
... ":class:`~lightning_app.core.flow.LightningFlow`"]
>>> from pprint import pprint
>>> pprint(_replace_imports_in_file(lns, {"app": "lightning_app"}))
['lightning.app',
'delete_cloud_lightning_apps',
'from lightning.app import',
'lightning_apps = []',
'lightning.app is ours',
'def _lightning_app():',
':class:`~lightning.app.core.flow.LightningFlow`']
"""
for n2, n1 in pkg_lut.items():
for i, ln in enumerate(lines):
lines[i] = re.sub(rf"([^_]|^){n1}([^_\w]|$)", rf"\1lightning.{n2}\2", ln)
return lines


# TODO: unify usage with assistant function, such that import this function in there
def copy_adjusted_modules(src_folder: str, pkg_name: str, lit_name: str, pkg_lut: dict) -> None:
"""Recursively replace imports in given folder."""
package_dir = os.path.join(src_folder, pkg_name)
all_files = glob.glob(os.path.join(package_dir, "**", "*.*"), recursive=True)
for fname in all_files:
local_path = fname.replace(package_dir + os.path.sep, "")
new_file = os.path.join(src_folder, "lightning", lit_name, local_path)
if not fname.endswith(".py"):
if not fname.endswith(".pyc"):
os.makedirs(os.path.dirname(new_file), exist_ok=True)
shutil.copy2(fname, new_file)
continue

with open(fname, encoding="utf-8") as fo:
py = fo.readlines()
py = _replace_imports_in_file(py, pkg_lut)
os.makedirs(os.path.dirname(new_file), exist_ok=True)
with open(new_file, "w", encoding="utf-8") as fo:
fo.writelines(py)


def create_mirror_package(src_folder: str, lit_pkg_mapping: dict) -> None:
"""Recursively replace imports in given folder.
>>> create_mirror_package(
... os.path.join(_PROJECT_ROOT, "src"),
... {"pytorch": "pytorch_lightning", "app": "lightning_app", "lite": "lightning_lite"}
... )
"""
mapping = lit_pkg_mapping.copy()
mapping.pop("lightning", None) # pop this key to avoid replacing `lightning` to `lightning.lightning`
for lit_name, pkg_name in mapping.items():
copy_adjusted_modules(src_folder, pkg_name, lit_name, mapping)


def _download_frontend(pkg_path: str):
"""Downloads an archive file for a specific release of the Lightning frontend and extracts it to the correct
directory."""
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/code-checks.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ on:
- "requirements/**"
- "src/**"
- "pyproject.toml" # includes mypy config
- "actions/**"
- ".actions/**"

concurrency:
group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref }}
Expand Down
1 change: 1 addition & 0 deletions MANIFEST.in
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,5 @@ exclude *.toml # project config
exclude requirements.txt
exclude __pycache__
include .actions/setup_tools.py
include .actions/assistant.py
include *.cff # citation info

0 comments on commit e33d09a

Please sign in to comment.