From 0cfb7df883a33c0362868ef7c89b2a098fa59d9a Mon Sep 17 00:00:00 2001 From: telamonian Date: Thu, 29 Aug 2024 20:18:19 -0400 Subject: [PATCH] refactored tarball creation/extraction to use `create_tarball`/`extract_tarball` --- comfy_cli/standalone.py | 68 ++++------------------ comfy_cli/utils.py | 124 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 129 insertions(+), 63 deletions(-) diff --git a/comfy_cli/standalone.py b/comfy_cli/standalone.py index d2b33f7..879fd06 100644 --- a/comfy_cli/standalone.py +++ b/comfy_cli/standalone.py @@ -1,17 +1,13 @@ import shutil import subprocess -import tarfile from pathlib import Path from typing import Optional import requests -from rich.live import Live -from rich.progress import Progress, TextColumn -from rich.table import Table from comfy_cli.constants import OS, PROC from comfy_cli.typing import PathLike -from comfy_cli.utils import download_progress, get_os, get_proc +from comfy_cli.utils import create_tarball, download_url, extract_tarball, get_os, get_proc from comfy_cli.uv import DependencyCompiler _here = Path(__file__).expanduser().resolve().parent @@ -36,6 +32,7 @@ def download_standalone_python( tag: str = "latest", flavor: str = "install_only", cwd: PathLike = ".", + show_progress: bool = True, ) -> PathLike: """grab a pre-built distro from the python-build-standalone project. See https://gregoryszorc.com/docs/python-build-standalone/main/""" @@ -60,7 +57,7 @@ def download_standalone_python( fname = f"{name}.tar.gz" url = f"{asset_url_prefix}/{fname}" - return download_progress(url, fname, cwd=cwd) + return download_url(url, fname, cwd=cwd, show_progress=show_progress) class StandalonePython: @@ -73,7 +70,8 @@ def FromDistro( flavor: str = "install_only", cwd: PathLike = ".", name: PathLike = "python", - ): + show_progress: bool = True, + ) -> "StandalonePython": fpath = download_standalone_python( platform=platform, proc=proc, @@ -81,28 +79,16 @@ def FromDistro( tag=tag, flavor=flavor, cwd=cwd, + show_progress=show_progress, ) return StandalonePython.FromTarball(fpath, name) @staticmethod - def FromTarball(fpath: PathLike, name: PathLike = "python") -> "StandalonePython": + def FromTarball(fpath: PathLike, name: PathLike = "python", show_progress: bool = True) -> "StandalonePython": fpath = Path(fpath) - - with tarfile.open(fpath) as tar: - info = tar.next() - old_name = info.name.split("/")[0] - - old_rpath = fpath.parent / old_name rpath = fpath.parent / name - # clean the tar file expand target and the final target - shutil.rmtree(old_rpath, ignore_errors=True) - shutil.rmtree(rpath, ignore_errors=True) - - with tarfile.open(fpath) as tar: - tar.extractall() - - shutil.move(old_rpath, rpath) + extract_tarball(inPath=fpath, outPath=rpath, show_progress=show_progress) return StandalonePython(rpath=rpath) @@ -177,40 +163,8 @@ def rehydrate_comfy_deps(self): ) self.dep_comp.install_wheels_directly() - def to_tarball(self, outPath: Optional[PathLike] = None, progress: bool = True): - outPath = self.rpath.with_suffix(".tgz") if outPath is None else Path(outPath) - - # do a little clean up prep - outPath.unlink(missing_ok=True) + def to_tarball(self, outPath: Optional[PathLike] = None, show_progress: bool = True): + # remove any __pycache__ before creating archive self.clean() - if progress: - fileSize = sum(f.stat().st_size for f in self.rpath.glob("**/*")) - - barProg = Progress() - addTar = barProg.add_task("[cyan]Creating tarball...", total=fileSize) - pathProg = Progress(TextColumn("{task.description}")) - pathTar = pathProg.add_task("") - - progress_table = Table.grid() - progress_table.add_row(barProg) - progress_table.add_row(pathProg) - - _size = 0 - - def _filter(tinfo: tarfile.TarInfo): - nonlocal _size - pathProg.update(pathTar, description=tinfo.path) - barProg.advance(addTar, _size) - _size = Path(tinfo.path).stat().st_size - return tinfo - else: - _filter = None - - with Live(progress_table, refresh_per_second=10): - with tarfile.open(outPath, "w:gz") as tar: - tar.add(self.rpath.relative_to(Path(".").expanduser().resolve()), filter=_filter) - - if progress: - barProg.advance(addTar, _size) - pathProg.update(pathTar, description="") + create_tarball(inPath=self.rpath, outPath=outPath, show_progress=show_progress) diff --git a/comfy_cli/utils.py b/comfy_cli/utils.py index 7dc9bcd..6a1e374 100644 --- a/comfy_cli/utils.py +++ b/comfy_cli/utils.py @@ -6,12 +6,16 @@ import platform import shutil import subprocess +import tarfile from pathlib import Path +from typing import Optional import psutil import requests import typer from rich import print, progress +from rich.live import Live +from rich.table import Table from comfy_cli.constants import DEFAULT_COMFY_WORKSPACE, OS, PROC from comfy_cli.typing import PathLike @@ -100,7 +104,13 @@ def f(incomplete: str) -> list[str]: return f -def download_progress(url: str, fname: PathLike, cwd: PathLike = ".", allow_redirects: bool = True) -> PathLike: +def download_url( + url: str, + fname: PathLike, + cwd: PathLike = ".", + allow_redirects: bool = True, + show_progress: bool = True, +) -> PathLike: """download url to local file fname and show a progress bar. See https://stackoverflow.com/q/37573483""" cwd = Path(cwd).expanduser().resolve() @@ -110,12 +120,114 @@ def download_progress(url: str, fname: PathLike, cwd: PathLike = ".", allow_redi if response.status_code != 200: response.raise_for_status() # Will only raise for 4xx codes, so... raise RuntimeError(f"Request to {url} returned status code {response.status_code}") - fsize = int(response.headers.get("Content-Length", 0)) - desc = "(Unknown total file size)" if fsize == 0 else "" response.raw.read = functools.partial(response.raw.read, decode_content=True) # Decompress if needed - with progress.wrap_file(response.raw, total=fsize, description=desc) as response_raw: - with fpath.open("wb") as f: - shutil.copyfileobj(response_raw, f) + with fpath.open("wb") as f: + if show_progress: + fsize = int(response.headers.get("Content-Length", 0)) + desc = f"downloading {fname}..." + ("(Unknown total file size)" if fsize == 0 else "") + + with progress.wrap_file(response.raw, total=fsize, description=desc) as response_raw: + shutil.copyfileobj(response_raw, f) + else: + shutil.copyfileobj(response.raw, f) return fpath + + +def extract_tarball( + inPath: PathLike, + outPath: Optional[PathLike] = None, + show_progress: bool = True, +): + inPath = Path(inPath).expanduser().resolve() + outPath = inPath.with_suffix("") if outPath is None else Path(outPath).expanduser().resolve() + + with tarfile.open(inPath) as tar: + info = tar.next() + old_name = info.name.split("/")[0] + # path to top-level of extraction result + extractPath = inPath.with_name(old_name) + + # clean both the extraction path and the final target path + shutil.rmtree(extractPath, ignore_errors=True) + shutil.rmtree(outPath, ignore_errors=True) + + if show_progress: + fileSize = inPath.stat().st_size + + barProg = progress.Progress() + barTask = barProg.add_task("[cyan]extracting tarball...", total=fileSize) + pathProg = progress.Progress(progress.TextColumn("{task.description}")) + pathTask = pathProg.add_task("") + + progress_table = Table.grid() + progress_table.add_row(barProg) + progress_table.add_row(pathProg) + + _size = 0 + + def _filter(tinfo: tarfile.TarInfo, _path: Optional[PathLike] = None): + nonlocal _size + pathProg.update(pathTask, description=tinfo.path) + barProg.advance(barTask, _size) + _size = tinfo.size + return tinfo + else: + _filter = None + + with Live(progress_table, refresh_per_second=10): + with tarfile.open(inPath) as tar: + tar.extractall(filter=_filter) + + if show_progress: + barProg.advance(barTask, _size) + pathProg.update(pathTask, description="") + + shutil.move(extractPath, outPath) + + +def create_tarball( + inPath: PathLike, + outPath: Optional[PathLike] = None, + cwd: Optional[PathLike] = None, + show_progress: bool = True, +): + cwd = Path("." if cwd is None else cwd).expanduser().resolve() + inPath = Path(inPath).expanduser().resolve() + outPath = inPath.with_suffix(".tgz") if outPath is None else Path(outPath).expanduser().resolve() + + # clean the archive target path + outPath.unlink(missing_ok=True) + + if show_progress: + fileSize = sum(f.stat().st_size for f in inPath.glob("**/*")) + + barProg = progress.Progress() + barTask = barProg.add_task("[cyan]creating tarball...", total=fileSize) + pathProg = progress.Progress(progress.TextColumn("{task.description}")) + pathTask = pathProg.add_task("") + + progress_table = Table.grid() + progress_table.add_row(barProg) + progress_table.add_row(pathProg) + + _size = 0 + + def _filter(tinfo: tarfile.TarInfo): + nonlocal _size + pathProg.update(pathTask, description=tinfo.path) + barProg.advance(barTask, _size) + _size = Path(tinfo.path).stat().st_size + return tinfo + else: + _filter = None + + with Live(progress_table, refresh_per_second=10): + with tarfile.open(outPath, "w:gz") as tar: + # don't include parent paths in archive + tar.add(inPath.relative_to(cwd), filter=_filter) + + if show_progress: + barProg.advance(barTask, _size) + pathProg.update(pathTask, description="")