Allow customizing the maximum allowed line drift (#90)

Previously it was hard-coded to 5 lines. Now we allow adjusting this number (this can be useful for custom codecs, such as invertible pyxl).

pyannotate_tools/annotations/__main__.py

@@ -24,6 +24,9 @@
                     help="Write output files")
 parser.add_argument('-j', '--processes', type=int, default=1, metavar="N",
                     help="Use N parallel processes (default no parallelism)")
+parser.add_argument('--max-line-drift', type=int, default=5, metavar="N",
+                    help="Maximum allowed line drift when inserting annotation"
+                         " (can be useful for custom codecs)")
 parser.add_argument('-v', '--verbose', action='store_true',
                     help="More verbose output")
 parser.add_argument('-q', '--quiet', action='store_true',
@@ -136,7 +139,8 @@ def main(args_override=None):
         nobackups=True,
         show_diffs=not args.quiet)
     if not rt.errors:
-        rt.refactor(args.files, write=args.write, num_processes=args.processes)
+        with FixAnnotateJson.max_line_drift_set(args.max_line_drift):
+            rt.refactor(args.files, write=args.write, num_processes=args.processes)
         if args.processes == 1:
             rt.summarize()
         else:

pyannotate_tools/fixes/fix_annotate_json.py

@@ -21,6 +21,7 @@
 import json  # noqa
 import os
 import re
+from contextlib import contextmanager
 
 from lib2to3.fixer_util import syms, touch_import
 from lib2to3.pgen2 import token
@@ -149,6 +150,7 @@ def count_args(node, results):
 class FixAnnotateJson(FixAnnotate):
 
     needed_imports = None
+    line_drift = 5
 
     def add_import(self, mod, name):
         if mod == self.current_module():
@@ -178,6 +180,16 @@ def make_annotation(self, node, results):
     # JSON data for the current file
     stub_json = None  # type: List[Dict[str, Any]]
 
+    @classmethod
+    @contextmanager
+    def max_line_drift_set(cls, max_drift):
+        old_drift = cls.line_drift
+        cls.line_drift = max_drift
+        try:
+            yield
+        finally:
+            cls.line_drift = old_drift
+
     @classmethod
     def init_stub_json_from_data(cls, data, filename):
         cls.stub_json = data
@@ -214,7 +226,7 @@ def get_annotation_from_stub(self, node, results, funcname):
             # If the line number is too far off, the source probably drifted
             # since the trace was collected; it's better to skip this node.
             # (Allow some drift, since decorators also cause an offset.)
-            if abs(node.get_lineno() - it['line']) >= 5:
+            if abs(node.get_lineno() - it['line']) >= self.line_drift:
                 self.log_message("%s:%d: '%s' signature from line %d too far away -- skipping" %
                                  (self.filename, node.get_lineno(), it['func_name'], it['line']))
                 return None

pyannotate_tools/fixes/tests/test_annotate_json_py2.py

@@ -410,6 +410,27 @@ def nop(a):
             """
         self.warns(a, a, "signature from line 10 too far away -- skipping", unchanged=True)
 
+    def test_line_number_drift_allowed(self):
+        self.setTestData(
+            [{"func_name": "yep",
+              "path": "<string>",
+              "line": 10,
+              "signature": {
+                  "arg_types": ["int"],
+                  "return_type": "int"},
+              }])
+        a = """\
+            def yep(a):
+                return a
+            """
+        b = """\
+            def yep(a):
+                # type: (int) -> int
+                return a
+            """
+        with FixAnnotateJson.max_line_drift_set(10):
+            self.check(a, b)
+
     def test_classmethod(self):
         # Class method names currently are returned without class name
         self.setTestData(