mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
93 lines
3.5 KiB
Python
93 lines
3.5 KiB
Python
import argparse
|
|
import base64
|
|
import copy
|
|
import csv
|
|
import hashlib
|
|
import io
|
|
import os
|
|
from pathlib import Path
|
|
from tempfile import NamedTemporaryFile
|
|
from typing import Final
|
|
from zipfile import ZipFile, ZipInfo
|
|
|
|
PACKAGE_DIR: Final = "litellm/cli"
|
|
PACKAGE_NAME: Final = "litellm_cli"
|
|
|
|
|
|
def editable_support(source: Path) -> tuple[tuple[str, bytes], ...]:
|
|
module: Final = "_litellm_cli_editable"
|
|
finder: Final = (
|
|
"import importlib.util\nimport sys\n\n"
|
|
"class _Finder:\n"
|
|
" @staticmethod\n"
|
|
" def find_spec(fullname, path=None, target=None):\n"
|
|
f" if fullname == {PACKAGE_NAME!r}:\n"
|
|
" return importlib.util.spec_from_file_location(\n"
|
|
f" fullname, {str(source / '__init__.py')!r},\n"
|
|
f" submodule_search_locations=[{str(source)!r}])\n"
|
|
" return None\n\n"
|
|
"sys.meta_path.insert(0, _Finder())\n"
|
|
)
|
|
return ((f"{module}.py", finder.encode()), (f"{module}.pth", f"import {module}\n".encode()))
|
|
|
|
|
|
def _name(name: str) -> str:
|
|
prefix: Final = f"{PACKAGE_DIR}/"
|
|
return f"{PACKAGE_NAME}/{name.removeprefix(prefix)}" if name.startswith(prefix) else name
|
|
|
|
|
|
def _info(original: ZipInfo) -> ZipInfo:
|
|
mapped: Final = copy.copy(original)
|
|
mapped.filename = _name(original.filename)
|
|
return mapped
|
|
|
|
|
|
def _record(entries: tuple[tuple[ZipInfo, bytes], ...], record: str) -> bytes:
|
|
buffer: Final = io.StringIO(newline="")
|
|
writer: Final = csv.writer(buffer, lineterminator="\n")
|
|
writer.writerows(
|
|
(
|
|
info.filename,
|
|
"sha256=" + base64.urlsafe_b64encode(hashlib.sha256(data).digest()).rstrip(b"=").decode(),
|
|
str(len(data)),
|
|
)
|
|
for info, data in entries
|
|
)
|
|
writer.writerow((record, "", ""))
|
|
return buffer.getvalue().encode()
|
|
|
|
|
|
def map_wheel(wheel: Path, *, source: Path | None = None) -> None:
|
|
with ZipFile(wheel) as archive:
|
|
records: Final = tuple(info for info in archive.infolist() if info.filename.endswith(".dist-info/RECORD"))
|
|
if len(records) != 1:
|
|
raise ValueError("A wheel must contain exactly one RECORD")
|
|
entries: Final = tuple(
|
|
(_info(info), archive.read(info))
|
|
for info in archive.infolist()
|
|
if info != records[0] and (source is None or not info.filename.startswith(f"{PACKAGE_DIR}/"))
|
|
)
|
|
additions: Final = () if source is None else tuple((ZipInfo(name), data) for name, data in editable_support(source))
|
|
contents: Final = (*entries, *additions)
|
|
names: Final = tuple(info.filename for info, _ in contents)
|
|
if len(names) != len(frozenset(names)):
|
|
raise ValueError("The CLI package mapping would overwrite a wheel member")
|
|
if source is None and not any(name.startswith(f"{PACKAGE_NAME}/") for name in names):
|
|
raise ValueError("The wheel does not contain the CLI package")
|
|
with NamedTemporaryFile(dir=wheel.parent, suffix=".whl", delete=False) as temporary:
|
|
destination: Final = Path(temporary.name)
|
|
try:
|
|
with ZipFile(destination, "w") as archive:
|
|
for info, data in contents:
|
|
archive.writestr(info, data)
|
|
archive.writestr(records[0], _record(contents, records[0].filename))
|
|
os.replace(destination, wheel)
|
|
finally:
|
|
destination.unlink(missing_ok=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser: Final = argparse.ArgumentParser()
|
|
parser.add_argument("wheels", nargs="+", type=Path)
|
|
for wheel in parser.parse_args().wheels:
|
|
map_wheel(wheel)
|