class Unnest(
resolver_op.ResolverOp,
arg_data_types=(resolver_op.DataTypes.ARTIFACT_MULTIMAP,),
return_data_type=resolver_op.DataTypes.ARTIFACT_MULTIMAP_LIST,
):
"""Unnest operator.
Unnest operator split a *`key` channel* of multiple artifacts into multiple
dicts each with a channel with single artifact. Pseudo code example:
Unnest({x: [x1, x2, x3]}, key=x)
= [{x: [x1]}, {x: [x2]}, {x: [x3]}]
For channels other than key channel remains the same. Pseudo code example:
Unnest({x: [x1, x2, x3], y: [y1]}, key=x)
= [{x: [x1], y: [y1]}, {x: [x2], y: [y1]}, {x: [x3], y: [y1]}]
"""
key = resolver_op.ResolverOpProperty(type=str)
def apply(self, input_dict: typing_utils.ArtifactMultiMap):
if self.key not in input_dict:
raise exceptions.FailedPreconditionError(
f'Input dict does not contain the key {self.key}. '
f'Available: {list(input_dict.keys())}')
main_channel = input_dict.get(self.key)
rest = {k: v for k, v in input_dict.items() if k != self.key}
result = []
for main_artifact in main_channel:
result.append({self.key: [main_artifact], **rest})
return result
class RepeatOp(resolver_op.ResolverOp):
num = resolver_op.ResolverOpProperty(type=int)
def __init__(self):
RepeatOp.last_created = self
def apply(self, input_dict):
return {key: value * self.num for key, value in input_dict.items()}
class BadProperty(resolver_op.ResolverOp): # pylint: disable=unused-variable
str_prop = resolver_op.ResolverOpProperty(type=str, default=42)
class Repeat(resolver_op.ResolverOp,
return_data_type=resolver_op.DataTypes.ARTIFACT_MULTIMAP_LIST):
n = resolver_op.ResolverOpProperty(type=int)
def apply(self, input_dict):
return [copy.deepcopy(input_dict) for _ in range(self.n)]
class Bar(resolver_op.ResolverOp):
bar = resolver_op.ResolverOpProperty(type=str, default='bar')
def apply(self, input_dict):
return input_dict
class Foo(resolver_op.ResolverOp):
foo = resolver_op.ResolverOpProperty(type=int)
def apply(self, input_dict):
return input_dict
class CustomOp(resolver_op.ResolverOp):
optional_mapping = resolver_op.ResolverOpProperty(
type=Optional[Mapping[str, str]], default=None)
def apply(self, input_dict):
return input_dict
class Bar(resolver_op.ResolverOp):
"""Dummy ResolverOperator Bar."""
bar = resolver_op.ResolverOpProperty(type=str)
def apply(self, input_dict):
return input_dict
class Foo(resolver_op.ResolverOp):
"""Dummy ResolverOperator Foo."""
foo = resolver_op.ResolverOpProperty(type=int)
def apply(self, input_dict):
return input_dict