###############################################################################
# (c) Copyright 2024 CERN for the benefit of the LHCb Collaboration #
# #
# This software is distributed under the terms of the GNU General Public #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING". #
# #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization #
# or submit itself to any jurisdiction. #
###############################################################################
import warnings
from copy import deepcopy
from typing import TypeVar, Union
from Functors.grammar import ComposedBoundFunctor, BoundFunctor # type: ignore[import]
from Functors.grammar import FunctorBase # type: ignore[import]
[docs]
Self_FunctorCollection = TypeVar("Self_FunctorCollection", bound="FunctorCollection")
[docs]
Type_FunctorDict = Union[Self_FunctorCollection, dict[str, Union[FunctorBase, str]]]
[docs]
class FunctorCollection:
"""
Class to hold the functor/variable dictionary.
Attributes:
functor_dict: dictionary of functors/variables, whose key is the field (branch) name suffix and value is the Functor code.
Methods are largely the same as for a dict class with specific requirements:
update : extends a given dictionary of functors to the collection, throwing a warning about overwriting an entry if it already exists.
pop : given a name or list of names it removes the functors from the collection, raising an error if the entry(ies) does(do) not exist.
add operator (+) : adds two FunctorCollection, for common entries warns the user about picking the entry from "base" FunctorCollection instance.
sub operator (-) : returns a FunctorCollection that contains unique entries between two FunctorCollections.
get_loki_functors(self): returns dictionary of loki functors
get_thor_functors(self): returns dictionary of thor functors
"""
[docs]
__slots__ = "functor_dict"
def __init__(self, functor_dict: Type_FunctorDict = {}) -> None:
"""
Default and copy constructor.
Args:
functor_dict (dict or FunctorCollection): input dictionary of functors/variables or a FunctorCollection instance.
For an input dict, keys are field (TTree branch) name suffixes and values are Functor code.
"""
# Allow for copy-construction
[docs]
functor_dict = (
deepcopy(functor_dict.functor_dict)
if isinstance(functor_dict, FunctorCollection)
else functor_dict
)
self.functor_dict: Type_FunctorDict = {}
try:
self.functor_dict = dict(functor_dict)
except ValueError:
raise ValueError(
f"FunctorCollection: Class must be initialised with a dict, not type {type(functor_dict)}. Please check!"
) from None
# Check that keys and value of expected type (replace this with typing module introduced in python 3.5 when available in newDV).
for k, v in self.functor_dict.items():
# checks for key
if not isinstance(k, str):
raise TypeError(
f"All keys for 'functor_dict' should be of type str; instead key {k} is of type {type(k)}. Please check!"
) from None
# checks for value
if not (
isinstance(v, str)
or isinstance(v, BoundFunctor)
or isinstance(v, ComposedBoundFunctor)
):
raise TypeError(
f"Functors have to be either of type string (for LoKi) or BoundFunctor (for ThOr) or ComposedBoundFunctor (for ThOr). The specified functor {v} is instead of type {type(v)}. Please check!"
) from None
[docs]
def update(self, functor_dict_new: Type_FunctorDict):
"""
Add new entries to the FunctorCollection given a dictionary.
Raises a user warning about overwriting an entry if it already exists.
"""
if isinstance(functor_dict_new, dict):
for key in functor_dict_new.keys():
if key in self.functor_dict:
warnings.warn(
f"FunctorCollection.update: The collection already contains an entry with same name {key}. Overwriting that entry."
)
return self.functor_dict.update(functor_dict_new)
else:
raise TypeError(
"FunctorCollection.update: Argument must be a dict. Please check!"
) from None
[docs]
def pop(self, functor_names: Union[str, list[str]]):
"""
Remove an entry or a list of entries from the FunctorCollection.
Args:
functor_names (str or list): entries to be removed.
Raises:
AttributeError: if the input is not adequate.
KeyError: if an entry to be removed does not exist.
"""
if isinstance(functor_names, list):
missingfunc = [
fname
for fname in functor_names
if fname not in self.functor_dict.keys()
]
if missingfunc:
raise KeyError(
f"FunctorCollection.pop: Failed attempt to remove non-existing key(s)\n{missingfunc}\nfrom the Collection. Nothing done."
) from None
else:
_ = [self.functor_dict.pop(fname) for fname in functor_names]
elif isinstance(functor_names, str):
self.pop([functor_names])
else:
raise AttributeError(
"FunctorCollection.pop: Argument must be a key or a list of keys. Please check!"
) from None
[docs]
def __repr__(self) -> str:
"""
Representation as a string in the common '<...>' format,
providing the number of ThOr and LoKi functors stored.
"""
n_thor_fctors, n_loki_fctors = len(self.get_thor_functors()), len(
self.get_loki_functors()
)
return f"<{self.__class__.__name__}: n_thor_fctors={n_thor_fctors}, n_loki_fctors={n_loki_fctors}>"
# Option str operator for printing
[docs]
def __str__(self) -> str:
"""
print(FunctorCollection) gives all the entries inside the FunctorCollection.
"""
return (
f"<FunctorCollection object at {hex(id(self))}:\n"
+ "".join(
[f" {key}: {val}\n" for key, val in (self.functor_dict).items()]
)
+ ">"
)
# define + operator
[docs]
def __add__(
self: Self_FunctorCollection, other: Self_FunctorCollection
) -> Self_FunctorCollection:
"""
Return all entries from two FunctorCollection e.g. FunctorCollection_1(A,B) + FuntorCollection_2(B,C) gives FunctorCollection(A,B,C) where B is from FunctorCollection_1.
Raises a user warning about common entries (this should rather be an error, I think)
"""
if not isinstance(other, FunctorCollection):
raise TypeError(
f"FunctorCollection.__add__: Input is not of type FunctorCollection, instead is of type {type(other)}. Please check!"
) from None
keys_self = list(self.functor_dict.keys())
keys_other = list(other.functor_dict.keys())
common_keys = list(set(keys_self).intersection(keys_other))
if common_keys:
# raise RuntimeError("The two functors have common entries {0}. Please check.".format(common_keys))
warnings.warn(
f"The two functors have common entries {common_keys}. For these we keep the entries from first collection."
)
tempdict = self.functor_dict.copy()
for key_other in keys_other:
if key_other in common_keys:
continue
tempdict[key_other] = other.functor_dict[key_other]
return self.__class__(tempdict)
# define += operator
[docs]
def __iadd__(
self: Self_FunctorCollection, other: Self_FunctorCollection
) -> Self_FunctorCollection:
"""
Return self adding the entries in FunctorCollection 'other'.
Raises:
TypeError: if the input type is not adequate, i.e a FunctorCollection instance.
"""
return self + other
[docs]
def __sub__(
self: Self_FunctorCollection, other: Self_FunctorCollection
) -> Self_FunctorCollection:
"""
Return unique entries from two FunctorCollections,
e.g. FunctorCollection_1(A,B) - FuntorCollection_2(B,C) gives FunctorCollection(A,C)
Raises:
RuntimeError: if the two functors do not have any unique entries.
"""
keys_self = list(self.functor_dict.keys())
keys_other = list(other.functor_dict.keys())
# Set of elements that are in either self or other, but not in their intersection
unique_keys = list(set(keys_self).symmetric_difference(set(keys_other)))
if len(unique_keys) == 0:
raise RuntimeError(
"The two functors do not have any unique entries. Please check"
) from None
tempdict = {}
for key in unique_keys:
if key in keys_other:
tempdict[key] = other.functor_dict[key]
else:
tempdict[key] = self.functor_dict[key]
return self.__class__(tempdict)
# define -= operator
[docs]
def __isub__(
self: Self_FunctorCollection, other: Self_FunctorCollection
) -> Self_FunctorCollection:
"""
Return self subtraction of the entries in FunctorCollection 'other'.
See __sub__ for a description of the behaviour.
"""
return self - other
[docs]
def __contains__(self, key: str) -> bool:
"""
Check for key containment.
Example:
>>> c = FunctorCollection({"PT": F.PT})
>>> "PT" in c
True
>>> "Nope" in c
False
"""
return key in self.functor_dict
[docs]
def get_loki_functors(self) -> dict[str, str]:
"""
Return dictionary of loki functors (checking if the values are strings or not).
"""
return {k: v for k, v in self.functor_dict.items() if isinstance(v, str)}
[docs]
def get_thor_functors(self) -> dict[str, FunctorBase]:
"""
Return dictionary of thor functors (checking if the values are not strings).
Explicitly type checking for ThOr needed here are they always of type 'Functors.grammar.BoundFunctor'(?)
"""
return {k: v for k, v in self.functor_dict.items() if not isinstance(v, str)}
[docs]
def __setitem__(self, key: str, value: Union[FunctorBase, str]) -> None:
"""
Add an item to the dictionary.
"""
if key in self.functor_dict:
warnings.warn(
f"FunctorCollection.__setitem__: The collection already contains the item with key {key}. Overwriting that entry."
)
self.functor_dict[key] = value
[docs]
def __getitem__(self, key: str) -> Union[FunctorBase, str]:
"""
Get an item to the dictionary.
"""
return self.functor_dict[key]
[docs]
def __eq__(self, other: Self_FunctorCollection) -> bool: # type: ignore[override]
"""
Equality purely on the contents of the collections,
i.e. the pairs (keys, values) of both dicts stored.
"""
return self.functor_dict == other.functor_dict
[docs]
def __len__(self) -> int:
"""
Implement len(FunctorCollection).
"""
return len(self.functor_dict)