mach-nix/mach_nix/data/nixpkgs.py
Tom Prince f101dc4c27 Update dependencies of python packages specified in packagesExtra in mkPython.
Otherwise, it is likely that there will be collisions building the python environment,
as the dependencies of packages in packagesExtra won't match with the dependencies
calculated by mach-nix.

As a side-effect, this also causes the dependencies of mach-nix built packages included
by the `nixpkgs` provider to be included in dependency resolution.
2022-01-27 15:42:42 -07:00

116 lines
4.4 KiB
Python

import json
from collections import UserDict
from dataclasses import dataclass
from typing import List
from mach_nix.cache import cached
from mach_nix.versions import parse_ver, Version
@dataclass
class NixpkgsPyPkg:
name: str
nix_key: str
ver: Version
class NixpkgsIndex(UserDict):
# mapping from pypi name to nix key
_aliases = dict(
torch='pytorch',
)
def __init__(self, nixpkgs_json_file, **kwargs):
with open(nixpkgs_json_file) as f:
data = json.load(f)
self.data = {}
self.requirements = {}
for nix_key, nix_data in data.items():
if nix_data is None:
continue
pname = nix_data["pname"]
version = parse_ver(nix_data["version"])
pname_key = pname.replace("_", "-").lower()
if pname_key not in self.data:
self.data[pname_key] = {}
if version not in self.data[pname_key]:
self.data[pname_key][version] = []
self.data[pname_key][version].append(nix_key)
if nix_data["requirements"] is not None:
self.requirements[nix_key] = nix_data["requirements"]
super(NixpkgsIndex, self).__init__(self.data, **kwargs)
def has_multiple_candidates(self, name):
count = 0
for nix_keys in self.data[name].values():
count += len(nix_keys)
if count > 1:
return True
return False
def get_all_candidates(self, name) -> List[NixpkgsPyPkg]:
result = []
for ver, nix_keys in self.data[name].items():
result += [NixpkgsPyPkg(name, nix_key, ver) for nix_key in nix_keys]
return result
def get_highest_ver(self, pkgs: List[NixpkgsPyPkg]):
# prioritizing similar names, reduces chance of infinite recursion
# (see problem with dateutil alias)
name_difference = lambda p: abs(len(p.name)-len(p.nix_key))
return max(pkgs, key=lambda p: (p.ver, -name_difference(p)))
@staticmethod
def is_same_ver(ver1, ver2, prefix_len):
"""
Check if the given versions share a common prefix of length `prefix_len`.
This ignores any epoch or pre/post-release versions. For the epoch,
it is likely that the nixpkg version doesn't include the epoch, and
just has the release version.
"""
def prefix(ver):
# Return a version prefix of the given length, padded with 0s.
return ver.release[:prefix_len] + (0,) * max(0, prefix_len - len(ver.release))
return prefix(ver1) == prefix(ver2)
@cached(lambda args: (args[1], args[2]))
def find_best_nixpkgs_candidate(self, name, ver):
"""
In case a python package has more than one candidate in nixpkgs
like `django` and `django_2_2`, this algo will select the right one.
"""
pkgs: List[NixpkgsPyPkg] = sorted(self.get_all_candidates(name), key=lambda pkg: pkg.ver)
if len(pkgs) == 1:
return pkgs[0].nix_key
# try to find nixpkgs candidate with closest version
remaining_pkgs = pkgs
for i in range(1, len(ver.release)):
same_ver = list(filter(lambda p: self.is_same_ver(ver, p.ver, i), remaining_pkgs))
if len(same_ver) == 1:
return same_ver[0].nix_key
elif len(same_ver) == 0:
# If there are no versions that match at this precision
# we pick the best version among the version that matched
# at the prior prefix.
break
remaining_pkgs = same_ver
# We've either fallen off the loop (in which case the versions match)
# or all remaining packages match the same length version prefix.
highest = self.get_highest_ver(remaining_pkgs).nix_key
print(f'Multiple nixpkgs attributes found for {name}-{ver}: {[p.nix_key for p in remaining_pkgs]}'
f"\nPicking '{highest}' as base attribute name.")
return highest
def exists(self, name, ver=None):
try:
candidates = [c for c in self.get_all_candidates(name)]
except KeyError:
return False
if ver:
return any(ver == c.ver for c in candidates)
return True
def get_requirements(self, name, ver):
nix_key = self.find_best_nixpkgs_candidate(name, ver)
return self.requirements.get(nix_key)