def get_kernel_locks()

in src/kernels/lockfile.py [0:0]

• 39 lines of code
• 13 McCabe index (conditional complexity)

def get_kernel_locks(repo_id: str, version_spec: str) -> KernelLock:
    """
    Get the locks for a kernel with the given version spec.

    The version specifier can be any valid Python version specifier:
    https://packaging.python.org/en/latest/specifications/version-specifiers/#version-specifiers
    """
    versions = _get_available_versions(repo_id)
    requirement = SpecifierSet(version_spec)
    accepted_versions = sorted(requirement.filter(versions.keys()))

    if len(accepted_versions) == 0:
        raise ValueError(
            f"No version of `{repo_id}` satisfies requirement: {version_spec}"
        )

    tag_for_newest = versions[accepted_versions[-1]]

    r = HfApi().repo_info(
        repo_id=repo_id, revision=tag_for_newest.target_commit, files_metadata=True
    )
    if r.sha is None:
        raise ValueError(
            f"Cannot get commit SHA for repo {repo_id} for tag {tag_for_newest.name}"
        )

    if r.siblings is None:
        raise ValueError(
            f"Cannot get sibling information for {repo_id} for tag {tag_for_newest.name}"
        )

    variant_files: Dict[str, List[Tuple[bytes, str]]] = {}
    for sibling in r.siblings:
        if sibling.rfilename.startswith("build/torch"):
            if sibling.blob_id is None:
                raise ValueError(f"Cannot get blob ID for {sibling.rfilename}")

            path = Path(sibling.rfilename)
            variant = path.parts[1]
            filename = Path(*path.parts[2:])

            hash = sibling.lfs.sha256 if sibling.lfs is not None else sibling.blob_id

            files = variant_files.setdefault(variant, [])

            # Encode as posix for consistent slash handling, then encode
            # as utf-8 for byte-wise sorting later.
            files.append((filename.as_posix().encode("utf-8"), hash))

    variant_locks = {}
    for variant, files in variant_files.items():
        m = hashlib.sha256()
        for filename_bytes, hash in sorted(files):
            # Filename as bytes.
            m.update(filename_bytes)
            # Git blob or LFS file hash as bytes.
            m.update(bytes.fromhex(hash))

        variant_locks[variant] = VariantLock(hash=f"sha256-{m.hexdigest()}")

    return KernelLock(repo_id=repo_id, sha=r.sha, variants=variant_locks)