Fix NotEqual position issue (Instagram#325)

Co-authored-by: Jimmy Lai <jimmylai@fb.com>
sk- · Jun 29, 2020 · 3a7ffaf · 3a7ffaf
1 parent 2d56ba7
commit 3a7ffaf
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 8 deletions.
diff --git a/libcst/_nodes/op.py b/libcst/_nodes/op.py
@@ -669,7 +669,7 @@ def _get_token(self) -> str:
 
 @add_slots
 @dataclass(frozen=True)
-class NotEqual(BaseCompOp):
+class NotEqual(BaseCompOp, _BaseOneTokenOp):
     """
     A comparison operator that can be used in a :class:`Comparison` expression.
 
@@ -691,7 +691,7 @@ def _validate(self) -> None:
         if self.value not in ["!=", "<>"]:
             raise CSTValidationError("Invalid value for NotEqual node.")
 
-    def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "BaseCompOp":
+    def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "NotEqual":
         return self.__class__(
             whitespace_before=visit_required(
                 self, "whitespace_before", self.whitespace_before, visitor
@@ -702,10 +702,8 @@ def _visit_and_replace_children(self, visitor: CSTVisitorT) -> "BaseCompOp":
             ),
         )
 
-    def _codegen_impl(self, state: CodegenState) -> None:
-        self.whitespace_before._codegen(state)
-        state.add_token(self.value)
-        self.whitespace_after._codegen(state)
+    def _get_token(self) -> str:
+        return self.value
 
 
 @add_slots

diff --git a/libcst/metadata/tests/test_position_provider.py b/libcst/metadata/tests/test_position_provider.py
@@ -9,7 +9,7 @@
 import libcst as cst
 from libcst import parse_module
 from libcst._batched_visitor import BatchableCSTVisitor
-from libcst._visitors import CSTTransformer
+from libcst._visitors import CSTVisitor
 from libcst.metadata import (
     CodeRange,
     MetadataWrapper,
@@ -38,7 +38,7 @@ def test_visitor_provider(self) -> None:
         """
         test = self
 
-        class DependentVisitor(CSTTransformer):
+        class DependentVisitor(CSTVisitor):
             METADATA_DEPENDENCIES = (PositionProvider,)
 
             def visit_Pass(self, node: cst.Pass) -> None:
@@ -49,6 +49,26 @@ def visit_Pass(self, node: cst.Pass) -> None:
         wrapper = MetadataWrapper(parse_module("pass"))
         wrapper.visit(DependentVisitor())
 
+    def test_equal_range(self) -> None:
+        test = self
+        expected_range = CodeRange((1, 4), (1, 6))
+
+        class EqualPositionVisitor(CSTVisitor):
+            METADATA_DEPENDENCIES = (PositionProvider,)
+
+            def visit_Equal(self, node: cst.Equal) -> None:
+                test.assertEqual(
+                    self.get_metadata(PositionProvider, node), expected_range
+                )
+
+            def visit_NotEqual(self, node: cst.NotEqual) -> None:
+                test.assertEqual(
+                    self.get_metadata(PositionProvider, node), expected_range
+                )
+
+        MetadataWrapper(parse_module("var == 1")).visit(EqualPositionVisitor())
+        MetadataWrapper(parse_module("var != 1")).visit(EqualPositionVisitor())
+
     def test_batchable_provider(self) -> None:
         test = self