Fix: tidy-imports fails to add missing imports when file has only local imports (#482)

* Add test for Quansight/deshaw#656

* Fix missing import not added.

This make sure the import get added globally and not to local import
block wrappers.

* ignore forwardref warnings

Author: Carreau

doc/conf.py

@@ -35,6 +35,10 @@
     "prompt_toolkit",
 ]
 
+suppress_warnings = [
+    "sphinx_autodoc_typehints.forward_reference",
+]
+
 html_theme_options = {
     'collapse_navigation': False,
     'navigation_depth': -1,

lib/python/pyflyby/_imports2s.py

@@ -737,14 +737,17 @@ def select_import_block_by_closest_prefix_match(self, imp, max_lineno):
           `SourceToSourceImportBlockTransformation`
         """
         # Create a data structure that annotates blocks with data by which
-        # we'll sort.
+        # we'll sort.  Only consider global import blocks, not local ones
+        # (wrapped in _LocalImportBlockWrapper), since new imports should
+        # always be added at the module level.
         annotated_blocks = [
             ( (max([0] + [len(imp.prefix_match(oimp))
                           for oimp in block.importset.imports]),
                block.input.endpos.lineno),
               block )
             for block in self.import_blocks
-            if block.input.endpos.lineno <= max_lineno+1 ]
+            if not isinstance(block, _LocalImportBlockWrapper)
+            and block.input.endpos.lineno <= max_lineno+1 ]
         if not annotated_blocks:
             raise NoImportBlockError()
         annotated_blocks.sort(key=lambda x: x[0])

pyproject.toml

@@ -43,6 +43,7 @@ py3 = "pyflyby._py:py_main"
 
 [project.optional-dependencies]
 test = [
+    'numpy',
     'beartype',
     'build',
     'coverage',

tests/test_imports2s_bis.py

@@ -227,6 +227,76 @@ def fun_2():
     assert output == expected
 
 
+@pytest.mark.parametrize("tidy_local_imports", (False, True))
+@pytest.mark.parametrize(
+    ("input_code", "expected_code"),
+    [
+        pytest.param(
+            """\
+def f():
+    import math
+    x = math.inf + inf
+""",
+            """\
+from numpy import inf
+
+def f():
+    import math
+    x = math.inf + inf
+""",
+            id="no_global_imports_with_local",
+        ),
+        pytest.param(
+            """\
+import os
+
+def f():
+    import math
+    x = math.inf + inf
+""",
+            """\
+import os
+from numpy import inf
+
+def f():
+    import math
+    x = math.inf + inf
+""",
+            id="existing_global_imports_with_local",
+        ),
+        pytest.param(
+            """\
+def f():
+    x = inf
+""",
+            """\
+from numpy import inf
+
+def f():
+    x = inf
+""",
+            id="no_global_no_local_imports",
+        ),
+    ],
+)
+def test_add_missing_imports_with_local_imports(
+    input_code, expected_code, tidy_local_imports
+):
+    """Regression test: missing imports must be added at module level even when
+    only local imports exist and tidy_local_imports is enabled.
+
+    """
+    input_block = PythonBlock(dedent(input_code).lstrip())
+    db = ImportDB("from numpy import inf")
+    output = fix_unused_and_missing_imports(
+        input_block, db=db, add_missing=True, remove_unused=False,
+        add_mandatory=False, tidy_local_imports=True,
+    )
+    result = str(output)
+    expected = dedent(expected_code).lstrip()
+    assert result == expected
+
+
 @pytest.mark.parametrize(
     "line,import_str,expected",
     [