def _handle_Index( # noqa
self, slice: cst.Index, node: cst.BaseExpression
) -> cst.BaseExpression:
value = slice.value
if isinstance(value, cst.Subscript):
new_slice = slice.with_changes(value=self._handle_Subscript(value))
return node.with_changes(slice=new_slice)
elif isinstance(value, cst.Attribute):
new_slice = slice.with_changes(value=self._add_annotation_to_imports(value))
return node.with_changes(slice=new_slice)
else:
return node
def _add_match_if_true(
oldtype: cst.BaseExpression, concrete_only_expr: cst.BaseExpression
) -> cst.BaseExpression:
"""
Given a BaseExpression in a type, add MatchIfTrue to it. This either
wraps in a Union type, or adds to the end of an existing Union type.
"""
if isinstance(oldtype, cst.Subscript) and oldtype.value.deep_equals(
cst.Name("Union")
):
# Add to the end of the value
return oldtype.with_changes(
slice=[*oldtype.slice, _get_match_if_true(concrete_only_expr)]
)
# Just wrap in a union type
return _get_wrapped_union_type(oldtype, _get_match_if_true(concrete_only_expr))