registry.py

# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from functools import partial
from types import FunctionType
from typing import Any, Dict, List, Optional, Union

from pytorch_lightning.utilities import rank_zero_info
from pytorch_lightning.utilities.exceptions import MisconfigurationException

_REGISTERED_FUNCTION = Dict[str, Any]


class FlashRegistry:
    """
    This class is used to register function or partial to a registry:

    Example::

        backbones = FlashRegistry("backbones")

        @backbones
        def my_model(nc_input=5, nc_output=6):
            return nn.Linear(nc_input, nc_output), nc_input, nc_output

        mlp, nc_input, nc_output = backbones("my_model")(nc_output=7)

        backbones(my_model, name="foo")
        assert backbones("foo")

    """

    def __init__(self, name: str, verbose: bool = False) -> None:
        self.name = name
        self.functions: List[_REGISTERED_FUNCTION] = []
        self._verbose = verbose

    def __len__(self) -> int:
        return len(self.functions)

    def __contains__(self, key) -> bool:
        return any(key == e["name"] for e in self.functions)

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}(name={self.name}, functions={self.functions})'

    def get(
        self,
        key: str,
        with_metadata: bool = False,
        strict: bool = True,
        **metadata,
    ) -> Union[callable, _REGISTERED_FUNCTION, List[_REGISTERED_FUNCTION], List[callable]]:
        """
        This function is used to gather matches from the registry:

        Args:
            key: Name of the registered function.
            with_metadata: Whether to include the associated metadata in the return value.
            strict: Whether to return all matches or just one.
            metadata: Metadata used to filter against existing registry item's metadata.
        """
        matches = [e for e in self.functions if key == e["name"]]
        if not matches:
            raise KeyError(f"Key: {key} is not in {repr(self)}")

        if metadata:
            matches = [m for m in matches if metadata.items() <= m["metadata"].items()]
            if not matches:
                raise KeyError("Found no matches that fit your metadata criteria. Try removing some metadata")

        matches = [e if with_metadata else e["fn"] for e in matches]
        return matches[0] if strict else matches

    def remove(self, key: str) -> None:
        self.functions = [f for f in self.functions if f["name"] != key]

    def _register_function(
        self,
        fn: callable,
        name: Optional[str] = None,
        override: bool = False,
        metadata: Optional[Dict[str, Any]] = None
    ):
        if not isinstance(fn, FunctionType) and not isinstance(fn, partial):
            raise MisconfigurationException(f"You can only register a function, found: {fn}")

        name = name or fn.__name__

        if self._verbose:
            rank_zero_info(f"Registering: {fn.__name__} function with name: {name} and metadata: {metadata}")

        item = {"fn": fn, "name": name, "metadata": metadata or {}}

        matching_index = self._find_matching_index(item)
        if override and matching_index is not None:
            self.functions[matching_index] = item
        else:
            if matching_index is not None:
                raise MisconfigurationException(
                    f"Function with name: {name} and metadata: {metadata} is already present within {self}."
                    " HINT: Use `override=True`."
                )
            self.functions.append(item)

    def _find_matching_index(self, item: _REGISTERED_FUNCTION) -> Optional[int]:
        for idx, fn in enumerate(self.functions):
            if (
                fn["fn"] == item["fn"] and fn["name"] == item["name"]
                and item["metadata"].items() <= fn["metadata"].items()
            ):
                return idx

    def __call__(
        self,
        fn: Optional[callable] = None,
        name: Optional[str] = None,
        override: bool = False,
        **metadata
    ) -> callable:
        """Register a function"""
        if fn is not None:
            self._register_function(fn=fn, name=name, override=override, metadata=metadata)
            return fn

        # raise the error ahead of time
        if not (name is None or isinstance(name, str)):
            raise TypeError(f'`name` must be a str, found {name}')

        def _register(cls):
            self._register_function(fn=cls, name=name, override=override, metadata=metadata)
            return cls

        return _register

    def available_keys(self) -> List[str]:
        return sorted(v["name"] for v in self.functions)