def leave_Subscript(
self, original_node: cst.Subscript, updated_node: cst.Subscript
) -> cst.Subscript:
if updated_node.value.deep_equals(cst.Name("Sequence")):
nodeslice = updated_node.slice
if isinstance(nodeslice, cst.Index):
possibleunion = nodeslice.value
if isinstance(possibleunion, cst.Subscript):
# Special case for Sequence[Union] so that we make more collapsed
# types.
if possibleunion.value.deep_equals(cst.Name("Union")):
return updated_node.with_changes(
slice=nodeslice.with_changes(
value=possibleunion.with_changes(
slice=[*possibleunion.slice, _get_do_not_care()]
)
)
)
# This is a sequence of some node, add DoNotCareSentinel here so that
# a person can add a do not care to a sequence that otherwise has
# valid matcher nodes.
return updated_node.with_changes(
slice=cst.Index(
_get_wrapped_union_type(nodeslice.value, _get_do_not_care())
)
)
raise Exception("Unexpected slice type for Sequence!")
return updated_node
def _handle_Index(self, slice: cst.Index, node: cst.Subscript) -> cst.Subscript:
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 _get_clean_type_from_subscript(
aliases: List[Alias], typecst: cst.Subscript) -> cst.BaseExpression:
if typecst.value.deep_equals(cst.Name("Sequence")):
# Lets attempt to widen the sequence type and alias it.
if len(typecst.slice) != 1:
raise Exception(
"Logic error, Sequence shouldn't have more than one param!")
inner_type = typecst.slice[0].slice
if not isinstance(inner_type, cst.Index):
raise Exception(
"Logic error, expecting Index for only Sequence element!")
inner_type = inner_type.value
if isinstance(inner_type, cst.Subscript):
clean_inner_type = _get_clean_type_from_subscript(
aliases, inner_type)
elif isinstance(inner_type, (cst.Name, cst.SimpleString)):
clean_inner_type = _get_clean_type_from_expression(
aliases, inner_type)
else:
raise Exception("Logic error, unexpected type in Sequence!")
return _get_wrapped_union_type(
typecst.deep_replace(inner_type, clean_inner_type),
_get_do_not_care(),
_get_match_if_true(typecst),
)
# We can modify this as-is to add our extra values
elif typecst.value.deep_equals(cst.Name("Union")):
return _get_clean_type_from_union(aliases, typecst)
else:
# Don't handle other types like "Literal", just widen them.
return _get_clean_type_from_expression(aliases, typecst)
def leave_Subscript(
self, original_node: cst.Subscript, updated_node: cst.Subscript
) -> cst.Subscript:
if updated_node.value.deep_equals(cst.Name("Union")):
# Take the original node, remove do not care so we have concrete types.
concrete_only_expr = _remove_do_not_care(original_node)
# Take the current subscript, add a MatchIfTrue node to it.
match_if_true_expr = _add_match_if_true(
_remove_do_not_care(updated_node), concrete_only_expr
)
return updated_node.with_changes(
slice=[
*updated_node.slice,
# Make sure that OneOf/AllOf types are widened to take all of the
# original SomeTypes, and also takes a MatchIfTrue, so that
# you can do something like OneOf(SomeType(), MatchIfTrue(lambda)).
# We could explicitly enforce that MatchIfTrue could not be used
# inside OneOf/AllOf clauses, but then if you want to mix and match you
# would have to use a recursive matches() inside your lambda which
# is super ugly.
cst.ExtSlice(cst.Index(_add_generic("OneOf", match_if_true_expr))),
cst.ExtSlice(cst.Index(_add_generic("AllOf", match_if_true_expr))),
# We use original node here, because we don't want MatchIfTrue
# to get modifications to child Union classes. If we allow
# that, we get MatchIfTrue nodes whose Callable takes in
# OneOf/AllOf and MatchIfTrue values, which is incorrect. MatchIfTrue
# only takes in cst nodes, and returns a boolean.
_get_match_if_true(concrete_only_expr),
]
)
return updated_node
def _handle_Subscript(self, node: cst.Subscript) -> cst.Subscript:
value = node.value
if isinstance(value, NAME_OR_ATTRIBUTE):
new_node = node.with_changes(
value=self._handle_NameOrAttribute(value))
else:
raise ValueError("Expected any indexed type to have")
if self._get_unique_qualified_name(node) in ("Type", "typing.Type"):
# Note: we are intentionally not handling qualification of
# anything inside `Type` because it's common to have nested
# classes, which we cannot currently distinguish from classes
# coming from other modules, appear here.
return new_node
slice = node.slice
if isinstance(slice, tuple):
new_slice = []
for item in slice:
value = item.slice.value
if isinstance(value, NAME_OR_ATTRIBUTE):
name = self._handle_NameOrAttribute(item.slice.value)
new_index = item.slice.with_changes(value=name)
new_slice.append(item.with_changes(slice=new_index))
else:
if isinstance(item.slice, cst.Index):
new_index = item.slice.with_changes(
value=self._handle_Index(item.slice, item))
item = item.with_changes(slice=new_index)
new_slice.append(item)
return new_node.with_changes(slice=tuple(new_slice))
elif isinstance(slice, cst.Index):
new_slice = self._handle_Index(slice)
return new_node.with_changes(slice=new_slice)
else:
return new_node
def _handle_Subscript(self, node: cst.Subscript) -> cst.Subscript:
slice = node.slice
if m.matches(node.value, m.Name(value="Type")):
return node
if isinstance(slice, list):
new_slice = []
for item in slice:
value = item.slice.value
if isinstance(value, cst.Attribute):
name = self._add_annotation_to_imports(item.slice.value)
new_index = item.slice.with_changes(value=name)
new_slice.append(item.with_changes(slice=new_index))
else:
if isinstance(item.slice, cst.Index) and not isinstance(
item.slice.value, cst.Name
):
new_index = item.slice.with_changes(
value=self._handle_Index(item.slice, item)
)
item = item.with_changes(slice=new_index, comma=None)
new_slice.append(item)
return node.with_changes(slice=new_slice)
elif isinstance(slice, cst.Index):
return self._handle_Index(slice, node)
else:
return node
def _get_clean_type_from_union(aliases: List[Alias],
typecst: cst.Subscript) -> cst.BaseExpression:
name = _get_alias_name(typecst)
value = typecst.with_changes(slice=[
*[_maybe_fix_sequence_in_union(aliases, slc) for slc in typecst.slice],
_get_match_metadata(),
_get_match_if_true(typecst),
])
return _wrap_clean_type(aliases, name, value)
def _get_clean_type_from_union(aliases: List[Alias],
typecst: cst.Subscript) -> cst.BaseExpression:
name = _get_alias_name(typecst)
value = typecst.with_changes(slice=[
# pyre-ignore We know .slice is a sequence. This can go away once we
# deprecate ExtSlice.
*[_maybe_fix_sequence_in_union(aliases, slc) for slc in typecst.slice],
_get_match_metadata(),
_get_match_if_true(typecst),
])
return _wrap_clean_type(aliases, name, value)
def _wrap_clean_type(aliases: List[Alias], name: Optional[str],
value: cst.Subscript) -> cst.BaseExpression:
if name is not None:
# We created an alias, lets use that, wrapping the alias in a do not care.
aliases.append(
Alias(name=name, type=cst.Module(body=()).code_for_node(value)))
return _get_wrapped_union_type(cst.Name(name), _get_do_not_care())
else:
# Couldn't name the alias, fall back to regular node creation, add do not
# care to the resulting type we widened.
return value.with_changes(slice=[*value.slice, _get_do_not_care()])
def _get_clean_type_from_subscript(
aliases: List[Alias], typecst: cst.Subscript) -> cst.BaseExpression:
# We can modify this as-is to add our extra values
if not typecst.value.deep_equals(cst.Name("Union")):
# Don't handle other types like "Literal", just widen them.
return _get_clean_type_from_expression(aliases, typecst)
name = _get_alias_name(typecst)
value = typecst.with_changes(slice=[
*typecst.slice,
_get_match_metadata(),
_get_match_if_true(typecst)
])
return _wrap_clean_type(aliases, name, value)
def leave_Subscript(self, original_node: cst.Subscript,
updated_node: cst.Subscript) -> cst.Subscript:
if original_node in self.in_match_if_true:
self.in_match_if_true.remove(original_node)
if original_node in self.fixup_nodes:
self.fixup_nodes.remove(original_node)
return updated_node.with_changes(slice=[
*updated_node.slice,
cst.SubscriptElement(
cst.Index(_add_generic("AtLeastN", original_node))),
cst.SubscriptElement(
cst.Index(_add_generic("AtMostN", original_node))),
])
return updated_node
def leave_Subscript(self, original_node: cst.Subscript,
updated_node: cst.Subscript) -> cst.Subscript:
if updated_node.value.deep_equals(cst.Name("Sequence")):
slc = updated_node.slice
# TODO: We can remove the instance check after ExtSlice is deprecated.
if not isinstance(slc, Sequence) or len(slc) != 1:
raise Exception(
"Unexpected number of sequence elements inside Sequence type "
+ "annotation!")
nodeslice = slc[0].slice
if isinstance(nodeslice, cst.Index):
possibleunion = nodeslice.value
if isinstance(possibleunion, cst.Subscript):
# Special case for Sequence[Union] so that we make more collapsed
# types.
if possibleunion.value.deep_equals(cst.Name("Union")):
return updated_node.with_deep_changes(
possibleunion,
slice=[
*possibleunion.slice,
_get_do_not_care(),
_get_match_metadata(),
],
)
# This is a sequence of some node, add DoNotCareSentinel here so that
# a person can add a do not care to a sequence that otherwise has
# valid matcher nodes.
return updated_node.with_changes(slice=(cst.SubscriptElement(
cst.Index(
_get_wrapped_union_type(
nodeslice.value,
_get_do_not_care(),
_get_match_metadata(),
))), ))
raise Exception("Unexpected slice type for Sequence!")
return updated_node
def leave_Subscript(self, original_node: cst.Subscript,
updated_node: cst.Subscript) -> cst.Subscript:
if updated_node.value.deep_equals(cst.Name("Union")):
# Take the original node, remove do not care so we have concrete types.
# Explicitly taking the original node because we want to discard nested
# changes.
concrete_only_expr = _remove_types(updated_node,
["DoNotCareSentinel"])
return updated_node.with_changes(slice=[
*updated_node.slice,
cst.SubscriptElement(
cst.Index(_add_generic("OneOf", concrete_only_expr))),
cst.SubscriptElement(
cst.Index(_add_generic("AllOf", concrete_only_expr))),
])
return updated_node