def __call__(self, do_fn_type: Type[beam.DoFn]) -> Type[beam.DoFn]:
"""Decorator which registers the given DoFn to the targeted models.
This decorator also installs type constraints on the DoFn to guard it
from invalid argument types.
Args:
do_fn_type: type(DoFn). The new audting DoFn class to decorate.
Returns:
type(DoFn). The decorated DoFn.
Raises:
TypeError. When the new type is not a DoFn.
"""
if not issubclass(do_fn_type, beam.DoFn):
raise TypeError('%r is not a subclass of DoFn' % do_fn_type)
# The "mro" (method resolution order) of a class is the list of types
# the class is derived from, including itself, in the order they are
# searched for methods and attributes.
# To learn more, see: https://stackoverflow.com/a/2010732/4859885.
base_types_of_do_fn_type = set(inspect.getmro(do_fn_type))
for kind in self._targeted_kinds:
registered_do_fn_types = self._DO_FN_TYPES_BY_KIND[kind]
if any(issubclass(r, do_fn_type) for r in registered_do_fn_types):
# Always keep the most-derived DoFn type.
continue
registered_do_fn_types -= base_types_of_do_fn_type
registered_do_fn_types.add(do_fn_type)
# Decorate the DoFn with type constraints that raise an error when args
# or return values have the wrong type.
with_input_types, with_output_types = (
typehints.with_input_types(
typehints.Union[self._targeted_model_types]),
typehints.with_output_types(base_validation_errors.BaseAuditError))
# TODO(#15613): The return type of functions with_input_types and
# with_output_types is Any, because these functions are not type
# annotated yet in Apache_beam library. Thus to return the appropriate
# type from function instead of Any. We used cast here.
return cast(Type[beam.DoFn],
with_input_types(with_output_types(do_fn_type)))
def __call__(self, do_fn):
"""Decorator which registers the given DoFn to the targeted models.
This decorator also installs type constraints on the DoFn to guard it
from invalid argument types.
Args:
do_fn: DoFn. The DoFn to decorate.
Returns:
do_fn. The decorated DoFn.
Raises:
TypeError. When the input argument is not a DoFn.
"""
if not issubclass(do_fn, beam.DoFn):
raise TypeError('%r is not a subclass of DoFn' % do_fn)
# The "mro" (method resolution order) of a class is the list of types
# the class is derived from, including itself, in the order they are
# searched for methods and attributes.
# To learn more see: https://stackoverflow.com/a/2010732/4859885.
base_classes_of_do_fn = set(inspect.getmro(do_fn))
for cls in self._model_classes:
registered_do_fns = self._DO_FNS_BY_MODEL_KIND[cls.__name__]
if any(issubclass(r, do_fn) for r in registered_do_fns):
# Always keep the most-derived DoFn.
continue
registered_do_fns -= base_classes_of_do_fn
registered_do_fns.add(do_fn)
# Decorate the DoFn with type constraints that raise an error when
# arguments or return values have the wrong type.
with_input_types, with_output_types = (
typehints.with_input_types(typehints.Union[self._model_classes]),
typehints.with_output_types(audit_errors.BaseAuditError))
return with_input_types(with_output_types(do_fn))
