Registries API¶

class Registry:
    name = None
    root = None
    pooch = None
    version = None
    priority: int = 0

    def __str__(self):
        return self.to_string(TextWrapper())

    def __hash__(self):
        return hash(str(self))

    def to_string(self, wrapper):
        from io import StringIO
        output = StringIO()
        self.to_file(output, wrapper)
        return output.getvalue()

#    def to_file(self, output, wrapper):
#        print(f'Registry<{self.name=},{self.root=},{self.pooch=},{self.version=},{self.priority=}>', file=output)

    def _inner_specification_parts(self, wrapper) -> list[str]:
        return [self.name, self.root, self.pooch, self.version]

    def specification_parts(self, wrapper=None) -> list[str]:
        return self._inner_specification_parts(wrapper if wrapper else TextWrapper())

    def specification_string(self) -> str:
        return ' '.join(self.specification_parts())

    def to_file(self, output, wrapper):
        print(wrapper.line('Registry:', self.specification_parts(wrapper)), file=output)

    def known(self, name: str, ext: str = None, strict: bool = False):
        pass

    def unique(self, name: str):
        pass

    def fullname(self, name: str, ext: str = None):
        pass

    def is_available(self, name: str, ext: str = None):
        pass

    def path(self, name: str, ext: str = None) -> Path:
        pass

    def filenames(self) -> list[str]:
        pass

    def search(self, regex: Pattern):
        """Return filenames containing the regex pattern, uses regex search"""
        regex = ensure_regex_pattern(regex)
        return [x for x in self.filenames() if regex.search(x) is not None]

    def match(self, regex: Pattern):
        """Return regex *matching* registered file names -- which *start* with the regex pattern"""
        regex = ensure_regex_pattern(regex)
        return [x for x in self.filenames() if regex.match(x) is not None]

    def contents(self, *args, **kwargs):
        """Return the text contents of a Registry file"""
        return self.path(*args, **kwargs).read_text()

class LocalRegistry(Registry):
    def __init__(self, name: str, root: str, priority: int = 10, recursive=True):
        self.name = name
        self.root = Path(root)
        self.version = mccode_antlr_version()
        self.priority = priority
        # A recursive registry is a whole searchable tree, as -I/--search-dir and
        # the configured component directories are. A non-recursive one holds only
        # the files directly in root -- what `mcstas` does for the working
        # directory, which it never descends into.
        self.recursive = _as_recursive(recursive)
        self._index = None  # lazy basename -> [paths] index of root's contents

    def __repr__(self):
        return (f'LocalRegistry({self.name!r}, {self.root!r}, {self.priority!r}, '
                f'recursive={self.recursive!r})')

    def __hash__(self):
        return hash(str(self))

    def file_contents(self):
        return {'name': self.name, 'root': self.root.as_posix(), 'priority': self.priority,
                'recursive': self.recursive}

    def _inner_specification_parts(self, wrapper=None) -> list[str]:
        parts = [self.name, wrapper.url(self.root.as_posix())]
        # Only non-default behaviour is spelled out, so specifications for the
        # usual recursive registry are unchanged.
        if not self.recursive:
            parts.append(NON_RECURSIVE_SPECIFICATION)
        return parts

    def _filetype_iterator(self, filetype: str):
        return self.root.glob(f'**/*.{filetype}' if self.recursive else f'*.{filetype}')

    def _file_index(self) -> dict[str, list[Path]]:
        # Every _file_iterator call is a full recursive walk of root (~0.2s for
        # mcstas-comps), and one translation makes many such calls (known/
        # fullname/path per data file, component lookups, ...). Walk once per
        # registry instance instead; the tree is static for a translation's
        # lifetime. Includes directories, matching glob('**/name') semantics.
        if self._index is None:
            import os
            index: dict[str, list[Path]] = {}
            if self.recursive:
                for dirpath, dirnames, filenames in os.walk(self.root):
                    base = Path(dirpath)
                    for n in dirnames + filenames:
                        index.setdefault(n, []).append(base / n)
            else:
                try:
                    names = os.listdir(self.root)
                except OSError:
                    names = []
                for n in names:
                    index[n] = [self.root / n]
            self._index = index
        return self._index

    def _file_iterator(self, name: str):
        # The index only answers plain basenames -- glob metacharacters or an
        # embedded path separator still need real glob matching.
        if any(c in name for c in '*?[') or '/' in name or '\\' in name:
            return self.root.glob(f'**/{name}' if self.recursive else name)
        return iter(self._file_index().get(name, []))

    def _exact_file_iterator(self, name: str):
        return self.root.glob(name)

    def known(self, name: str, ext: str = None, strict: bool = False):
        compare = _name_plus_suffix(name, ext)
        return len(list(self._file_iterator(compare))) > 0

    def unique(self, name: str):
        return len(list(self._file_iterator(name))) == 1

    def fullname(self, name: str, ext: str = None, exact: bool = False):
        compare = _name_plus_suffix(name, ext)
        # Candidate sets that had several entries which _dedupe_identical_paths
        # refused to collapse, i.e. real ambiguities. A later stage can legitimately
        # find nothing at all, so without remembering these an ambiguous lookup
        # ends up reported as "No match" -- the opposite of what happened.
        ambiguous: list[list[Path]] = []

        def resolve(candidates: list[Path]):
            if len(candidates) == 1:
                return candidates[0]
            if len(candidates) > 1:
                if (deduped := _dedupe_identical_paths(candidates)) is not None:
                    return deduped
                ambiguous.append(candidates)
            return None

        # Complete match, then a complete match if name is missing the extension
        for candidates in (self._exact_file_iterator(compare), self._exact_file_iterator(name)):
            if (found := resolve(list(candidates))) is not None:
                return found
        if not exact:
            for candidates in (self._file_iterator(compare), self._file_iterator(name)):
                if (found := resolve(list(candidates))) is not None:
                    return found
        # Or matching *any* file that contains name
        if (found := resolve(list(self._file_iterator(name)))) is not None:
            return found
        if ambiguous:
            distinct = sorted({p for candidates in ambiguous for p in candidates}, key=str)
            listing = '\n'.join(f'  {p}' for p in distinct)
            raise RuntimeError(
                f'Ambiguous match for {compare} under {self.root}: {len(distinct)} '
                f'candidates with differing contents\n{listing}\n'
                'Narrow the search with -I/--search-dir, or remove the duplicates.'
            )
        raise RuntimeError(f'No match for {compare} or {name} under {self.root}')

    def is_available(self, name: str, ext: str = None):
        return self.known(name, ext)

    def path(self, name: str, ext: str = None, exact: bool = False) -> Path:
        return self.root.joinpath(self.fullname(name, ext, exact))

    def filenames(self) -> list[str]:
        return [str(x) for x in self.root.glob('**' if self.recursive else '*')]

    def __eq__(self, other):
        if not isinstance(other, Registry):
            return False
        if other.name != self.name:
            return False
        if other.root != self.root:
            return False
        if getattr(other, 'recursive', True) != self.recursive:
            return False
        return True

class GitHubRegistry(RemoteRegistry):
    def __init__(self, name: str, url: str, version: str, filename: str | None = None,
                 registry: str | dict | None = None, priority: int = 0):

        if filename is None:
            filename = f'{name}-registry.txt'
        super().__init__(name, url, version, filename, priority)

        # If registry is a string url, we expect the registry file to be available from _that_ url
        self._stashed_registry = None
        if isinstance(registry, str) and simple_url_validator(registry, file_ok=True):
            self._stashed_registry = registry
        self._registry_dict = registry if isinstance(registry, dict) else None

    def _build_pooch(self):
        from os import access, R_OK, W_OK
        registry = self._registry_dict
        if registry is None and self._stashed_registry:
            registry = f'{self._stashed_registry}/raw/{self.version}/'

        safe_name = _safe_cache_component(self.name, "unnamed")
        safe_version = _safe_cache_component(self.version, "unversioned")
        registry_file = self.filename or 'pooch-registry.txt'
        safe_file = _safe_cache_component(registry_file, 'pooch_registry.txt')

        base_url = f'{self.url}/raw/{self.version}/'
        cache_path = pooch.os_cache(f'mccodeantlr/{safe_name}')
        registry_file_path = cache_path.joinpath(safe_version, safe_file)
        if registry_file_path.exists() and registry_file_path.is_file() and access(registry_file_path, R_OK):
            with registry_file_path.open('r') as file:
                registry = {k: v for k, v in [x.strip().split(maxsplit=1) for x in file.readlines() if len(x)]}
        else:
            # We allow a full-dictionary to be provided, otherwise we expect the registry file to be available from the
            # base_url where all subsequent files are also expected to be available
            if not isinstance(registry, dict):
                r = _fetch_registry_with_retry((registry or base_url) + registry_file)
                if not r.ok:
                    raise RuntimeError(f"Could not retrieve {r.url} because {r.reason}")
                registry = {k: v for k, v in [x.split(maxsplit=1) for x in r.text.split('\n') if len(x)]}
            # stash-away the registry file to be re-read next time
            check = registry_file_path.parent
            last = Path('/')
            while not check.exists() and check != last:
                last, check = check, check.parent
            # check is now a directory that exists, it may be the root of the filesystem
            if access(check, W_OK):
                registry_file_path.parent.mkdir(parents=True, exist_ok=True)
                with registry_file_path.open('w') as file:
                    file.writelines('\n'.join([f'{k} {v}' for k, v in registry.items()]))
            else:
                logger.warning(f'Can not output {registry_file_path}, you lack write permissions for {check}')

        version = self.version if self.version and self.version.startswith('v') else None

        return pooch.create(
            path=cache_path,
            base_url=base_url,
            version=version,
            version_dev="main",
            registry=registry,
        )

    @property
    def registry(self):
        return self._stashed_registry

    def _inner_specification_parts(self, wrapper):
        filename = self.filename or f'{self.name}-registry.txt'
        items = [self.name, wrapper.url(self.url or ''), self.version or '', filename]
        if self._stashed_registry:
            items.append(wrapper.url(self._stashed_registry))
        return items


    def file_contents(self) -> dict[str, str]:
        fc = super().file_contents()
        fc['registry'] = self._stashed_registry or ''
        return fc

    @classmethod
    def file_keys(cls) -> tuple[str, ...]:
        return super().file_keys() + ('registry',)

class InMemoryRegistry(Registry):
    def __init__(self, name, priority: int = 100, files=None, origins=None, **components):
        self.name = name
        self.version = mccode_antlr_version()
        self.priority = priority
        self.files: dict[str, bytes] = {}
        for key, value in (files or {}).items():
            self.files[key] = _decode_stored_bytes(value)
        self.origins: dict[str, str] = dict(origins or {})
        for key, value in components.items():
            self.add_comp(key, value)
        self._materialized: set[str] = set()

    @property
    def components(self) -> dict[str, str]:
        out = {}
        for key, value in self.files.items():
            try:
                out[key] = value.decode('utf-8')
            except UnicodeDecodeError:
                continue
        return out

    def add(self, name: str, definition, origin: str | None = None):
        self.files[name] = definition.encode('utf-8') if isinstance(definition, str) else bytes(definition)
        if origin is not None:
            self.origins[name] = origin

    def add_comp(self, name: str, definition, origin: str | None = None):
        if not name.lower().endswith('.comp'):
            name += '.comp'
        self.add(name, definition, origin=origin)

    def add_instr(self, name: str, definition, origin: str | None = None):
        if not name.lower().endswith('.instr'):
            name += '.instr'
        self.add(name, definition, origin=origin)

    @classmethod
    def file_keys(cls) -> tuple[str, ...]:
        return 'name', 'priority', 'files', 'origins'

    def file_contents(self) -> dict:
        from base64 import b64encode
        return {
            'name': self.name,
            'priority': self.priority,
            'files': {k: b64encode(v).decode('ascii') for k, v in self.files.items()},
            'origins': dict(self.origins),
        }

    def filenames(self) -> list[str]:
        return list(self.files.keys())

    def fullname(self, name: str, ext: str | None = None):
        full_name = name if ext is None else name + ext
        return full_name if full_name in self.files else None

    def known(self, name: str, ext: str | None = None, strict: bool = False):
        return self.fullname(name, ext=ext) is not None

    def is_available(self, name: str, ext: str | None = None):
        return self.known(name, ext)

    def unique(self, name: str):
        return sum(1 for key in self.files if name in key) == 1

    def contents_bytes(self, name: str, ext: str | None = None) -> bytes:
        full_name = self.fullname(name, ext=ext)
        if full_name is None:
            err_name = name if ext is None else name + ext
            raise KeyError(f'InMemoryRegistry does not know of {err_name}')
        return self.files[full_name]

    def contents(self, name: str, ext: str | None = None):
        raw = self.contents_bytes(name, ext)
        try:
            return raw.decode('utf-8')
        except UnicodeDecodeError as error:
            full = name if ext is None else name + ext
            raise RuntimeError(f'{full} in registry {self.name!r} is not UTF-8 text; '
                               'use contents_bytes or path to read it') from error

    @property
    def root(self) -> Path:
        """Directory the entries occupy once written out.

        Content-addressed, so repeated translations and concurrent processes
        converge on the same directory.
        """
        digest = hashlib.sha256()
        for key in sorted(self.files):
            digest.update(key.encode('utf-8'))
            digest.update(hashlib.sha256(self.files[key]).digest())
        return self._cache_root() / self.name / digest.hexdigest()[:16]

    @staticmethod
    @cache
    def _cache_root() -> Path:
        base = Path(pooch.os_cache('mccodeantlr')) / 'in-memory'
        try:
            base.mkdir(parents=True, exist_ok=True)
            return base
        except OSError:
            from tempfile import mkdtemp
            return Path(mkdtemp(prefix='mccodeantlr-in-memory-'))

    def _ensure_materialized(self, key: str) -> Path:
        """Write the one entry to disk so it can be opened by path."""
        target = self.root / key
        if key in self._materialized:
            return target
        payload = self.files[key]
        if not (target.exists() and target.read_bytes() == payload):
            from os import getpid
            target.parent.mkdir(parents=True, exist_ok=True)
            # Write-then-rename so a concurrent reader never sees a partial file.
            tmp = target.with_name(f'{target.name}.{getpid()}.tmp')
            tmp.write_bytes(payload)
            tmp.replace(target)
            self._materialized.add(key)
        return target

    def path(self, name: str, ext: str | None = None) -> Path:
        full_name = self.fullname(name, ext=ext)
        if full_name is None:
            err_name = name if ext is None else name + ext
            raise KeyError(f'InMemoryRegistry does not know of {err_name!r}')
        return self._ensure_materialized(full_name)

    def to_file(self, output, wrapper):
        print(
            wrapper.line('InMemoryRegistry:', [self.name, f'({len(self.files)} files)']),
            file=output
        )

    def __eq__(self, other):
        if not isinstance(other, InMemoryRegistry):
            return False
        return self.name == other.name and self.files == other.files

    def __hash__(self):
        return hash(
            (self.name, tuple(sorted(
                 (k, hashlib.sha256(v).hexdigest()) for k, v in self.files.items()
            )))
        )