def _visit_preorder(
type_spec: computation_types.Type,
fn: Callable[[computation_types.Type, T], T],
context: T,
):
context = fn(type_spec, context)
for child_type in type_spec.children():
visit_preorder(child_type, fn, context)
def contains(type_signature: computation_types.Type,
predicate: _TypePredicate) -> bool:
"""Checks if `type_signature` contains any types that pass `predicate`."""
if predicate(type_signature):
return True
for child in type_signature.children():
if contains(child, predicate):
return True
return False
def count(type_signature: computation_types.Type,
predicate: _TypePredicate) -> int:
"""Returns the number of types in `type_signature` matching `predicate`.
Args:
type_signature: A tree of `computation_type.Type`s to count.
predicate: A Python function that takes a type as a parameter and returns a
boolean value.
"""
counter = 1 if predicate(type_signature) else 0
counter += sum(count(child, predicate) for child in type_signature.children())
return counter
def visit_preorder(type_signature: computation_types.Type,
fn: Callable[[computation_types.Type, T], T], context: T):
"""Recursively calls `fn` on the possibly nested structure `type_signature`.
Walks the tree in a preorder manner. Updates `context` on the way down with
the appropriate information, as defined in `fn`.
Args:
type_signature: A `computation_types.Type`.
fn: A function to apply to each of the constituent elements of
`type_signature` with the argument `context`. Must return an updated
version of `context` which incorporated the information we'd like to track
as we move down the type tree.
context: Initial state of information to be passed down the tree.
"""
context = fn(type_signature, context)
for child_type in type_signature.children():
visit_preorder(child_type, fn, context)
def _preorder_types(type_signature: computation_types.Type):
yield type_signature
for child in type_signature.children():
yield from _preorder_types(child)