Allow FlattenSentinel in matcher decorator return annotations

libcst/codegen/gen_type_mapping.py

@@ -19,6 +19,7 @@
 generated_code.append("# This file was generated by libcst.codegen.gen_type_mapping")
 generated_code.append("from typing import Dict as TypingDict, Type, Union")
 generated_code.append("")
+generated_code.append("from libcst._flatten_sentinel import FlattenSentinel")
 generated_code.append("from libcst._maybe_sentinel import MaybeSentinel")
 generated_code.append("from libcst._removal_sentinel import RemovalSentinel")
 generated_code.append("from libcst._nodes.base import CSTNode")
@@ -47,12 +48,10 @@
     base_uses = nodeuses[nodebases[node]]
     if node_uses.maybe or base_uses.maybe:
         valid_return_types.append("MaybeSentinel")
-    if (
-        node_uses.optional
-        or node_uses.sequence
-        or base_uses.optional
-        or base_uses.sequence
-    ):
+    if node_uses.sequence or base_uses.sequence:
+        valid_return_types.append(f'FlattenSentinel["{nodebases[node].__name__}"]')
+        valid_return_types.append("RemovalSentinel")
+    elif node_uses.optional or base_uses.optional:
         valid_return_types.append("RemovalSentinel")
     generated_code.append(f'    {name}: Union[{", ".join(valid_return_types)}],')
 generated_code.append("}")

libcst/matchers/_return_types.py

@@ -7,6 +7,7 @@
 # This file was generated by libcst.codegen.gen_type_mapping
 from typing import Dict as TypingDict, Type, Union
 
+from libcst._flatten_sentinel import FlattenSentinel
 from libcst._maybe_sentinel import MaybeSentinel
 from libcst._nodes.base import CSTNode
 from libcst._nodes.expression import (
@@ -214,17 +215,25 @@
     Add: BaseBinaryOp,
     AddAssign: BaseAugOp,
     And: BaseBooleanOp,
-    AnnAssign: Union[BaseSmallStatement, RemovalSentinel],
+    AnnAssign: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     Annotation: Annotation,
-    Arg: Union[Arg, RemovalSentinel],
+    Arg: Union[Arg, FlattenSentinel["Arg"], RemovalSentinel],
     AsName: AsName,
-    Assert: Union[BaseSmallStatement, RemovalSentinel],
-    Assign: Union[BaseSmallStatement, RemovalSentinel],
+    Assert: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
+    Assign: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     AssignEqual: Union[AssignEqual, MaybeSentinel],
-    AssignTarget: Union[AssignTarget, RemovalSentinel],
+    AssignTarget: Union[AssignTarget, FlattenSentinel["AssignTarget"], RemovalSentinel],
     Asynchronous: Asynchronous,
     Attribute: BaseExpression,
-    AugAssign: Union[BaseSmallStatement, RemovalSentinel],
+    AugAssign: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     Await: BaseExpression,
     BinaryOperation: BaseExpression,
     BitAnd: BaseBinaryOp,
@@ -235,54 +244,86 @@
     BitXor: BaseBinaryOp,
     BitXorAssign: BaseAugOp,
     BooleanOperation: BaseExpression,
-    Break: Union[BaseSmallStatement, RemovalSentinel],
+    Break: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     Call: BaseExpression,
-    ClassDef: Union[BaseStatement, RemovalSentinel],
+    ClassDef: Union[BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel],
     Colon: Union[Colon, MaybeSentinel],
     Comma: Union[Comma, MaybeSentinel],
     Comment: Comment,
     CompFor: CompFor,
     CompIf: CompIf,
     Comparison: BaseExpression,
-    ComparisonTarget: Union[ComparisonTarget, RemovalSentinel],
+    ComparisonTarget: Union[
+        ComparisonTarget, FlattenSentinel["ComparisonTarget"], RemovalSentinel
+    ],
     ConcatenatedString: BaseExpression,
-    Continue: Union[BaseSmallStatement, RemovalSentinel],
-    Decorator: Union[Decorator, RemovalSentinel],
-    Del: Union[BaseSmallStatement, RemovalSentinel],
+    Continue: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
+    Decorator: Union[Decorator, FlattenSentinel["Decorator"], RemovalSentinel],
+    Del: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     Dict: BaseExpression,
     DictComp: BaseExpression,
-    DictElement: Union[BaseDictElement, RemovalSentinel],
+    DictElement: Union[
+        BaseDictElement, FlattenSentinel["BaseDictElement"], RemovalSentinel
+    ],
     Divide: BaseBinaryOp,
     DivideAssign: BaseAugOp,
-    Dot: Union[Dot, RemovalSentinel],
-    Element: Union[BaseElement, RemovalSentinel],
+    Dot: Union[Dot, FlattenSentinel["Dot"], RemovalSentinel],
+    Element: Union[BaseElement, FlattenSentinel["BaseElement"], RemovalSentinel],
     Ellipsis: BaseExpression,
     Else: Else,
-    EmptyLine: Union[EmptyLine, RemovalSentinel],
+    EmptyLine: Union[EmptyLine, FlattenSentinel["EmptyLine"], RemovalSentinel],
     Equal: BaseCompOp,
-    ExceptHandler: Union[ExceptHandler, RemovalSentinel],
-    ExceptStarHandler: Union[ExceptStarHandler, RemovalSentinel],
-    Expr: Union[BaseSmallStatement, RemovalSentinel],
+    ExceptHandler: Union[
+        ExceptHandler, FlattenSentinel["ExceptHandler"], RemovalSentinel
+    ],
+    ExceptStarHandler: Union[
+        ExceptStarHandler, FlattenSentinel["ExceptStarHandler"], RemovalSentinel
+    ],
+    Expr: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     Finally: Finally,
     Float: BaseExpression,
     FloorDivide: BaseBinaryOp,
     FloorDivideAssign: BaseAugOp,
-    For: Union[BaseStatement, RemovalSentinel],
+    For: Union[BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel],
     FormattedString: BaseExpression,
-    FormattedStringExpression: Union[BaseFormattedStringContent, RemovalSentinel],
-    FormattedStringText: Union[BaseFormattedStringContent, RemovalSentinel],
+    FormattedStringExpression: Union[
+        BaseFormattedStringContent,
+        FlattenSentinel["BaseFormattedStringContent"],
+        RemovalSentinel,
+    ],
+    FormattedStringText: Union[
+        BaseFormattedStringContent,
+        FlattenSentinel["BaseFormattedStringContent"],
+        RemovalSentinel,
+    ],
     From: From,
-    FunctionDef: Union[BaseStatement, RemovalSentinel],
+    FunctionDef: Union[
+        BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel
+    ],
     GeneratorExp: BaseExpression,
-    Global: Union[BaseSmallStatement, RemovalSentinel],
+    Global: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     GreaterThan: BaseCompOp,
     GreaterThanEqual: BaseCompOp,
-    If: Union[BaseStatement, RemovalSentinel],
+    If: Union[BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel],
     IfExp: BaseExpression,
     Imaginary: BaseExpression,
-    Import: Union[BaseSmallStatement, RemovalSentinel],
-    ImportAlias: Union[ImportAlias, RemovalSentinel],
-    ImportFrom: Union[BaseSmallStatement, RemovalSentinel],
+    Import: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
+    ImportAlias: Union[ImportAlias, FlattenSentinel["ImportAlias"], RemovalSentinel],
+    ImportFrom: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     ImportStar: ImportStar,
     In: BaseCompOp,
     IndentedBlock: BaseSuite,
@@ -291,30 +332,44 @@
     Is: BaseCompOp,
     IsNot: BaseCompOp,
     Lambda: BaseExpression,
-    LazyImport: Union[BaseSmallStatement, RemovalSentinel],
-    LazyImportFrom: Union[BaseSmallStatement, RemovalSentinel],
+    LazyImport: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
+    LazyImportFrom: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     LeftCurlyBrace: LeftCurlyBrace,
-    LeftParen: Union[LeftParen, MaybeSentinel, RemovalSentinel],
+    LeftParen: Union[
+        LeftParen, MaybeSentinel, FlattenSentinel["LeftParen"], RemovalSentinel
+    ],
     LeftShift: BaseBinaryOp,
     LeftShiftAssign: BaseAugOp,
     LeftSquareBracket: LeftSquareBracket,
     LessThan: BaseCompOp,
     LessThanEqual: BaseCompOp,
     List: BaseExpression,
     ListComp: BaseExpression,
-    Match: Union[BaseStatement, RemovalSentinel],
+    Match: Union[BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel],
     MatchAs: MatchPattern,
     MatchCase: MatchCase,
     MatchClass: MatchPattern,
-    MatchKeywordElement: Union[MatchKeywordElement, RemovalSentinel],
+    MatchKeywordElement: Union[
+        MatchKeywordElement, FlattenSentinel["MatchKeywordElement"], RemovalSentinel
+    ],
     MatchList: MatchPattern,
     MatchMapping: MatchPattern,
-    MatchMappingElement: Union[MatchMappingElement, RemovalSentinel],
+    MatchMappingElement: Union[
+        MatchMappingElement, FlattenSentinel["MatchMappingElement"], RemovalSentinel
+    ],
     MatchOr: MatchPattern,
-    MatchOrElement: Union[MatchOrElement, RemovalSentinel],
+    MatchOrElement: Union[
+        MatchOrElement, FlattenSentinel["MatchOrElement"], RemovalSentinel
+    ],
     MatchPattern: MatchPattern,
     MatchSequence: MatchPattern,
-    MatchSequenceElement: Union[MatchSequenceElement, RemovalSentinel],
+    MatchSequenceElement: Union[
+        MatchSequenceElement, FlattenSentinel["MatchSequenceElement"], RemovalSentinel
+    ],
     MatchSingleton: MatchPattern,
     MatchStar: MatchStar,
     MatchTuple: MatchPattern,
@@ -328,61 +383,87 @@
     Multiply: BaseBinaryOp,
     MultiplyAssign: BaseAugOp,
     Name: BaseExpression,
-    NameItem: Union[NameItem, RemovalSentinel],
+    NameItem: Union[NameItem, FlattenSentinel["NameItem"], RemovalSentinel],
     NamedExpr: BaseExpression,
     Newline: Newline,
-    Nonlocal: Union[BaseSmallStatement, RemovalSentinel],
+    Nonlocal: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     Not: BaseUnaryOp,
     NotEqual: BaseCompOp,
     NotIn: BaseCompOp,
     Or: BaseBooleanOp,
-    Param: Union[Param, MaybeSentinel, RemovalSentinel],
+    Param: Union[Param, MaybeSentinel, FlattenSentinel["Param"], RemovalSentinel],
     ParamSlash: Union[ParamSlash, MaybeSentinel],
     ParamSpec: ParamSpec,
     ParamStar: Union[ParamStar, MaybeSentinel],
     Parameters: Parameters,
     ParenthesizedWhitespace: Union[BaseParenthesizableWhitespace, MaybeSentinel],
-    Pass: Union[BaseSmallStatement, RemovalSentinel],
+    Pass: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     Plus: BaseUnaryOp,
     Power: BaseBinaryOp,
     PowerAssign: BaseAugOp,
-    Raise: Union[BaseSmallStatement, RemovalSentinel],
-    Return: Union[BaseSmallStatement, RemovalSentinel],
+    Raise: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
+    Return: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
     RightCurlyBrace: RightCurlyBrace,
-    RightParen: Union[RightParen, MaybeSentinel, RemovalSentinel],
+    RightParen: Union[
+        RightParen, MaybeSentinel, FlattenSentinel["RightParen"], RemovalSentinel
+    ],
     RightShift: BaseBinaryOp,
     RightShiftAssign: BaseAugOp,
     RightSquareBracket: RightSquareBracket,
     Semicolon: Union[Semicolon, MaybeSentinel],
     Set: BaseExpression,
     SetComp: BaseExpression,
-    SimpleStatementLine: Union[BaseStatement, RemovalSentinel],
+    SimpleStatementLine: Union[
+        BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel
+    ],
     SimpleStatementSuite: BaseSuite,
     SimpleString: BaseExpression,
     SimpleWhitespace: Union[BaseParenthesizableWhitespace, MaybeSentinel],
     Slice: BaseSlice,
     StarredDictComp: BaseExpression,
-    StarredDictElement: Union[BaseDictElement, RemovalSentinel],
+    StarredDictElement: Union[
+        BaseDictElement, FlattenSentinel["BaseDictElement"], RemovalSentinel
+    ],
     StarredElement: BaseExpression,
     Subscript: BaseExpression,
-    SubscriptElement: Union[SubscriptElement, RemovalSentinel],
+    SubscriptElement: Union[
+        SubscriptElement, FlattenSentinel["SubscriptElement"], RemovalSentinel
+    ],
     Subtract: BaseBinaryOp,
     SubtractAssign: BaseAugOp,
     TemplatedString: BaseExpression,
-    TemplatedStringExpression: Union[BaseTemplatedStringContent, RemovalSentinel],
-    TemplatedStringText: Union[BaseTemplatedStringContent, RemovalSentinel],
+    TemplatedStringExpression: Union[
+        BaseTemplatedStringContent,
+        FlattenSentinel["BaseTemplatedStringContent"],
+        RemovalSentinel,
+    ],
+    TemplatedStringText: Union[
+        BaseTemplatedStringContent,
+        FlattenSentinel["BaseTemplatedStringContent"],
+        RemovalSentinel,
+    ],
     TrailingWhitespace: TrailingWhitespace,
-    Try: Union[BaseStatement, RemovalSentinel],
-    TryStar: Union[BaseStatement, RemovalSentinel],
+    Try: Union[BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel],
+    TryStar: Union[BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel],
     Tuple: BaseExpression,
-    TypeAlias: Union[BaseSmallStatement, RemovalSentinel],
-    TypeParam: Union[TypeParam, RemovalSentinel],
+    TypeAlias: Union[
+        BaseSmallStatement, FlattenSentinel["BaseSmallStatement"], RemovalSentinel
+    ],
+    TypeParam: Union[TypeParam, FlattenSentinel["TypeParam"], RemovalSentinel],
     TypeParameters: TypeParameters,
     TypeVar: TypeVar,
     TypeVarTuple: TypeVarTuple,
     UnaryOperation: BaseExpression,
-    While: Union[BaseStatement, RemovalSentinel],
-    With: Union[BaseStatement, RemovalSentinel],
-    WithItem: Union[WithItem, RemovalSentinel],
+    While: Union[BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel],
+    With: Union[BaseStatement, FlattenSentinel["BaseStatement"], RemovalSentinel],
+    WithItem: Union[WithItem, FlattenSentinel["WithItem"], RemovalSentinel],
     Yield: BaseExpression,
 }

libcst/matchers/_visitors.py

@@ -99,11 +99,23 @@ def _annotation_is_union(annotation: object) -> bool:
     )
 
 
+def _unwrap_generic_alias(annotation: object) -> Type[object]:
+    # A parameterized generic such as FlattenSentinel[BaseStatement] is not a
+    # class, so issubclass() would raise TypeError on it. Fall back to the
+    # generic's origin class, which is what we actually want to compare.
+    origin = getattr(annotation, "__origin__", None)
+    if isinstance(origin, type):
+        return origin
+    return cast(Type[object], annotation)
+
+
 def _get_possible_annotated_classes(annotation: object) -> List[Type[object]]:
     if _annotation_is_union(annotation):
-        return getattr(annotation, "__args__", [])
+        return [
+            _unwrap_generic_alias(arg) for arg in getattr(annotation, "__args__", [])
+        ]
     else:
-        return [cast(Type[object], annotation)]
+        return [_unwrap_generic_alias(annotation)]
 
 
 def _get_valid_leave_annotations_for_classes(

libcst/matchers/tests/test_visitors.py

@@ -160,6 +160,28 @@ def _string_visit(
         # Instantiating this class should not raise any errors
         TestVisitor()
 
+    def test_valid_transformer_leave_return_flatten(self) -> None:
+        class TestVisitor(MatcherDecoratableTransformer):
+            @leave(m.FunctionDef())
+            def _function_def_leave(
+                self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
+            ) -> Union[cst.BaseStatement, cst.FlattenSentinel[cst.BaseStatement]]:
+                return updated_node
+
+        # Instantiating this class should not raise any errors
+        TestVisitor()
+
+    def test_valid_transformer_leave_return_bare_flatten(self) -> None:
+        class TestVisitor(MatcherDecoratableTransformer):
+            @leave(m.FunctionDef())
+            def _function_def_leave(
+                self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
+            ) -> Union[cst.BaseStatement, cst.FlattenSentinel]:
+                return updated_node
+
+        # Instantiating this class should not raise any errors
+        TestVisitor()
+
     def test_invalid_collector_visit_return(self) -> None:
         class TestVisitor(MatcherDecoratableVisitor):
             @visit(m.SimpleString())