text: # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import pickle
  from typing import Union
  
  import libcst as cst
  import libcst.matchers as m
  from libcst.matchers import (
      leave,
      MatchDecoratorMismatch,
      MatcherDecoratableTransformer,
      MatcherDecoratableVisitor,
      visit,
  )
  from libcst.testing.utils import UnitTest
  
  
  class MatchersVisitLeaveDecoratorTypingTest(UnitTest):
      def test_valid_collector_simple(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @visit(m.SimpleString())
              def _string_visit(self, node: cst.SimpleString) -> None:
                  pass
  
              @leave(m.SimpleString())
              def _string_leave(self, original_node: cst.SimpleString) -> None:
                  pass
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_transformer_simple(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @visit(m.SimpleString())
              def _string_visit(self, node: cst.SimpleString) -> None:
                  pass
  
              @leave(m.SimpleString())
              def _string_leave(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> cst.SimpleString:
                  return updated_node
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_transformer_base_class(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString())
              def _string_leave(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> cst.BaseExpression:
                  return updated_node
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_collector_visit_union(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @visit(m.SimpleString() | m.Name())
              def _string_visit(self, node: Union[cst.SimpleString, cst.Name]) -> None:
                  pass
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_transformer_visit_union(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @visit(m.SimpleString() | m.Name())
              def _string_visit(self, node: Union[cst.SimpleString, cst.Name]) -> None:
                  pass
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_collector_visit_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @visit(m.SimpleString() | m.Name())
              def _string_visit(self, node: cst.BaseExpression) -> None:
                  pass
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_transformer_visit_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @visit(m.SimpleString() | m.Name())
              def _string_visit(self, node: cst.BaseExpression) -> None:
                  pass
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_collector_leave_union(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @leave(m.SimpleString() | m.Name())
              def _string_leave(self, node: Union[cst.SimpleString, cst.Name]) -> None:
                  pass
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_transformer_leave_union(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString() | m.Name())
              def _string_leave(
                  self,
                  original_node: Union[cst.SimpleString, cst.Name],
                  updated_node: Union[cst.SimpleString, cst.Name],
              ) -> Union[cst.SimpleString, cst.Name]:
                  return updated_node
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_collector_leave_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @leave(m.SimpleString() | m.Name())
              def _string_leave(self, node: cst.BaseExpression) -> None:
                  pass
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_transformer_leave_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString() | m.Name())
              def _string_leave(
                  self,
                  original_node: cst.BaseExpression,
                  updated_node: cst.BaseExpression,
              ) -> cst.BaseExpression:
                  return updated_node
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_transformer_leave_return_maybe(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.AssignEqual())
              def _assign_equal_leave(
                  self, original_node: cst.AssignEqual, updated_node: cst.AssignEqual
              ) -> Union[cst.AssignEqual, cst.MaybeSentinel]:
                  return updated_node
  
          # Instantiating this class should not raise any errors
          TestVisitor()
  
      def test_valid_transformer_leave_return_remove(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.AssignTarget())
              def _string_visit(
                  self, original_node: cst.AssignTarget, updated_node: cst.AssignTarget
              ) -> Union[cst.AssignTarget, cst.RemovalSentinel]:
                  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())
              def _string_visit(self, node: cst.SimpleString) -> bool:
                  return False
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit should only decorate functions that do not return",
          ):
              TestVisitor()
  
      def test_invalid_transformer_visit_return(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @visit(m.SimpleString())
              def _string_visit(self, node: cst.SimpleString) -> bool:
                  return False
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit should only decorate functions that do not return",
          ):
              TestVisitor()
  
      def test_invalid_transformer_visit_num_params(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @visit(m.SimpleString())
              def _string_visit(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit should decorate functions which take 1 parameter",
          ):
              TestVisitor()
  
      def test_invalid_collector_visit_num_params(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @visit(m.SimpleString())
              def _string_visit(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit should decorate functions which take 1 parameter",
          ):
              TestVisitor()
  
      def test_invalid_transformer_leave_num_params(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString())
              def _string_leave(
                  self, original_node: cst.SimpleString
              ) -> cst.SimpleString:
                  return original_node
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave should decorate functions which take 2 parameters",
          ):
              TestVisitor()
  
      def test_invalid_collector_leave_num_params(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @leave(m.SimpleString())
              def _string_leave(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave should decorate functions which take 1 parameter",
          ):
              TestVisitor()
  
      def test_invalid_collector_leave_return(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @leave(m.SimpleString())
              def _string_leave(self, original_node: cst.SimpleString) -> bool:
                  return False
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave should only decorate functions that do not return",
          ):
              TestVisitor()
  
      def test_invalid_transformer_leave_return_invalid_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString())
              def _string_visit(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> cst.BaseParenthesizableWhitespace:
                  return cst.SimpleWhitespace("")
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave decorated function cannot return the type BaseParenthesizableWhitespace",
          ):
              TestVisitor()
  
      def test_invalid_transformer_leave_return_wrong_type(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString())
              def _string_visit(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> cst.Pass:
                  return cst.Pass()
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave decorated function cannot return the type Pass",
          ):
              TestVisitor()
  
      def test_invalid_transformer_leave_return_invalid_maybe(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString())
              def _string_visit(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> Union[cst.SimpleString, cst.MaybeSentinel]:
                  return updated_node
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave decorated function cannot return the type MaybeSentinel",
          ):
              TestVisitor()
  
      def test_invalid_transformer_leave_return_invalid_remove(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString())
              def _string_visit(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> Union[cst.SimpleString, cst.RemovalSentinel]:
                  return updated_node
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave decorated function cannot return the type RemovalSentinel",
          ):
              TestVisitor()
  
      def test_invalid_transformer_leave_return_invalid_union(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString() | m.Name())
              def _string_leave(
                  self,
                  original_node: Union[cst.SimpleString, cst.Name],
                  updated_node: Union[cst.SimpleString, cst.Name],
              ) -> Union[cst.SimpleString, cst.Pass]:
                  return cst.SimpleString('""')
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave decorated function cannot return the type Pass",
          ):
              TestVisitor()
  
      def test_invalid_collector_visit_union(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @visit(m.SimpleString() | m.Name())
              def _string_visit(self, node: cst.SimpleString) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit can be called with Name but the decorated function parameter annotations do not include this type",
          ):
              TestVisitor()
  
      def test_invalid_transformer_visit_union(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @visit(m.SimpleString() | m.Name())
              def _string_visit(self, node: cst.SimpleString) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit can be called with Name but the decorated function parameter annotations do not include this type",
          ):
              TestVisitor()
  
      def test_invalid_collector_visit_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @visit(m.SimpleString() | m.Pass())
              def _string_visit(self, node: cst.BaseExpression) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit can be called with Pass but the decorated function parameter annotations do not include this type",
          ):
              TestVisitor()
  
      def test_invalid_transformer_visit_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @visit(m.SimpleString() | m.Pass())
              def _string_visit(self, node: cst.BaseExpression) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit can be called with Pass but the decorated function parameter annotations do not include this type",
          ):
              TestVisitor()
  
      def test_invalid_collector_leave_union(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @leave(m.SimpleString() | m.Name())
              def _string_leave(self, node: cst.SimpleString) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave can be called with Name but the decorated function parameter annotations do not include this type",
          ):
              TestVisitor()
  
      def test_invalid_transformer_leave_union(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString() | m.Name())
              def _string_leave(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> cst.BaseExpression:
                  return updated_node
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave can be called with Name but the decorated function parameter annotations do not include this type",
          ):
              TestVisitor()
  
      def test_invalid_collector_leave_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @leave(m.SimpleString() | m.Pass())
              def _string_leave(self, node: cst.BaseExpression) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave can be called with Pass but the decorated function parameter annotations do not include this type",
          ):
              TestVisitor()
  
      def test_invalid_transformer_leave_superclass(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString() | m.Pass())
              def _string_leave(
                  self,
                  original_node: cst.BaseExpression,
                  updated_node: cst.BaseExpression,
              ) -> cst.BaseExpression:
                  return updated_node
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave can be called with Pass but the decorated function parameter annotations do not include this type",
          ):
              TestVisitor()
  
      def test_bad_visit_collecter_decorator(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @visit(m.SimpleString())
              def visit_SimpleString(self, node: cst.SimpleString) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit should not decorate functions that are concrete visit or leave methods",
          ):
              TestVisitor()
  
      def test_bad_leave_collecter_decorator(self) -> None:
          class TestVisitor(MatcherDecoratableVisitor):
              @leave(m.SimpleString())
              def leave_SimpleString(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave should not decorate functions that are concrete visit or leave methods",
          ):
              TestVisitor()
  
      def test_bad_visit_transform_decorator(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @visit(m.SimpleString())
              def visit_SimpleString(self, node: cst.SimpleString) -> None:
                  pass
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@visit should not decorate functions that are concrete visit or leave methods",
          ):
              TestVisitor()
  
      def test_bad_leave_transform_decorator(self) -> None:
          class TestVisitor(MatcherDecoratableTransformer):
              @leave(m.SimpleString())
              def leave_SimpleString(
                  self, original_node: cst.SimpleString, updated_node: cst.SimpleString
              ) -> cst.SimpleString:
                  return updated_node
  
          # Instantiating this class should raise a runtime error
          with self.assertRaisesRegex(
              MatchDecoratorMismatch,
              "@leave should not decorate functions that are concrete visit or leave methods",
          ):
              TestVisitor()
  
      def test_pickleable_exception(self) -> None:
          original = MatchDecoratorMismatch("func", "message")
          serialized = pickle.dumps(original)
          unserialized = pickle.loads(serialized)
          self.assertEqual(original.message, unserialized.message)
          self.assertEqual(original.func, unserialized.func)
  