Number-replacer using offset list rather than dict.
Also using explicit unary positive for clarity.
@@ -6,11 +6,12 @@ import ast
from .operator import Operator
-# Maps mutation count to number offsets
-OFFSET_MAP = {
- 0: 1,
- 1: -1,
-}
+# List of offsets that we apply to numbers in the AST. Each index into the list
+# corresponds to single mutation.
+OFFSETS = [
+ +1,
+ -1,
+]
class NumberReplacer(Operator):
@@ -18,13 +19,13 @@ class NumberReplacer(Operator):
def visit_Num(self, node): # noqa # pylint: disable=invalid-name
"Visit a number node."
- return self.visit_mutation_site(node, len(OFFSET_MAP))
+ return self.visit_mutation_site(node, len(OFFSETS))
def mutate(self, node, idx):
"""Modify the numeric value on `node`."""
- assert idx in OFFSET_MAP, 'received count with no associated offset'
+ assert idx < len(OFFSETS), 'received count with no associated offset'
- offset = OFFSET_MAP[idx]
+ offset = OFFSETS[idx]
new_node = ast.Num(n=node.n + offset)
return new_node
