Merge pull request #91 from Cheukting/nested_with

src/shed/_codemods.py

@@ -9,7 +9,7 @@
 import re
 from ast import literal_eval
 from functools import wraps
-from typing import List, Tuple
+from typing import List, Tuple, Union
 
 import libcst as cst
 import libcst.matchers as m
@@ -70,7 +70,7 @@ def _run_codemods(code: str, min_version: Tuple[int, int]) -> str:
 
     if imports_hypothesis(code):  # pragma: no cover
         mod = attempt_hypothesis_codemods(context, mod)
-    mod = ShedFixers(context).transform_module(mod)
+    mod = ShedFixers(context, min_version).transform_module(mod)
     return mod.code
 
 
@@ -137,6 +137,10 @@ class ShedFixers(VisitorBasedCodemodCommand):
 
     DESCRIPTION = "Fix a variety of style, performance, and correctness issues."
 
+    def __init__(self, context, min_version):
+        super().__init__(context)
+        self.min_version = min_version
+
     @m.call_if_inside(m.Raise(exc=m.Name(value="NotImplemented")))
     def leave_Name(self, _, updated_node):  # noqa
         return updated_node.with_changes(value="NotImplementedError")
@@ -676,3 +680,66 @@ def _remove_recursive_helper(cls, bool_node):
                     **{side: side_node.args[0]},
                 )
         return bool_node
+
+    # rewrite nested `with` statement - code source from https://github.com/lensvol/pybetter/blob/master/pybetter/transformers/nested_withs.py
+
+    @leave(m.With())
+    def remove_nested_with(self, _, updated_node):
+        if self.min_version < (3, 9):
+            return updated_node
+
+        candidate_with: cst.With = updated_node
+        compound_items: List[cst.WithItem] = []
+        final_body: cst.BaseSuite = candidate_with.body
+
+        def has_leading_comment(node: Union[cst.SimpleStatementLine, cst.With]) -> bool:
+            return any([line.comment is not None for line in node.leading_lines])
+
+        header = m.AllOf(
+            m.TrailingWhitespace(),
+            m.MatchIfTrue(lambda h: h.comment is not None),
+        )
+        footer = [m.ZeroOrMore(), m.EmptyLine(comment=m.Comment()), m.ZeroOrMore()]
+
+        def has_footer_comment(body):
+            return m.matches(body, m.IndentedBlock(footer=footer))
+
+        while not (
+            # There is no way to meaningfully represent comments inside
+            # multi-line `with` statements due to how Python grammar is
+            # written, so we do not try to transform such `with` statements
+            # lest we lose something important in the comments.
+            has_leading_comment(candidate_with)
+            or m.matches(candidate_with.body, m.IndentedBlock(header=header))
+            # There is no meaningful way `async with` can be merged into
+            # the compound `with` statement.
+            or candidate_with.asynchronous
+        ):
+            compound_items.extend(candidate_with.items)
+            final_body = candidate_with.body
+
+            if not (
+                isinstance(final_body.body[0], cst.With) and len(final_body.body) == 1
+            ):
+                break  # pragma: no cover  # only reachable on some Python versions
+
+            candidate_with = cst.ensure_type(final_body.body[0], cst.With)
+
+        if len(compound_items) <= 1:
+            return updated_node
+
+        final_body = cst.ensure_type(final_body, cst.IndentedBlock)
+        topmost_body = cst.ensure_type(updated_node.body, cst.IndentedBlock)
+
+        if has_footer_comment(topmost_body) and not has_footer_comment(final_body):
+            final_body = final_body.with_changes(
+                footer=(*final_body.footer, *topmost_body.footer)
+            )
+
+        return updated_node.with_changes(
+            body=final_body,
+            items=compound_items,
+            # Black will only format with parens if they're there to start, so:
+            lpar=cst.LeftParen(),
+            rpar=cst.RightParen(),
+        )

tests/recorded/nested_with_38.txt

@@ -0,0 +1,49 @@
+# test nested with statement NOT get reformatted in 3.8 or below
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+        pass
+
+with make_context_manager(1) as cm1, make_context_manager(2) as cm2:
+    with make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2, make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1, make_context_manager(2) as cm2:
+    with make_context_manager(2) as cm2, make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+        with make_context_manager(3) as cm3:
+            with make_context_manager(4) as cm4:
+                pass
+
+================================================================================
+
+# test nested with statement NOT get reformatted in 3.8 or below
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+        pass
+
+with make_context_manager(1) as cm1, make_context_manager(2) as cm2:
+    with make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2, make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1, make_context_manager(2) as cm2:
+    with make_context_manager(2) as cm2, make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+        with make_context_manager(3) as cm3:
+            with make_context_manager(4) as cm4:
+                pass

tests/recorded/nested_with_39.txt

@@ -0,0 +1,62 @@
+# test nested with statement get reformatted in 3.9 or above
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+        pass
+    # Preserve this comment
+
+with make_context_manager(1) as cm1, make_context_manager(2) as cm2:
+    with make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2, make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1, make_context_manager(2) as cm2:
+    with make_context_manager(2) as cm2, make_context_manager(3) as cm3:
+        pass
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+        with make_context_manager(3) as cm3:
+            with make_context_manager(4) as cm4:
+                pass
+
+================================================================================
+
+# test nested with statement get reformatted in 3.9 or above
+
+with make_context_manager(1) as cm1, make_context_manager(2) as cm2:
+    pass
+    # Preserve this comment
+
+with (
+    make_context_manager(1) as cm1,
+    make_context_manager(2) as cm2,
+    make_context_manager(3) as cm3,
+):
+    pass
+
+with (
+    make_context_manager(1) as cm1,
+    make_context_manager(2) as cm2,
+    make_context_manager(3) as cm3,
+):
+    pass
+
+with (
+    make_context_manager(1) as cm1,
+    make_context_manager(2) as cm2,
+    make_context_manager(2) as cm2,
+    make_context_manager(3) as cm3,
+):
+    pass
+
+with (
+    make_context_manager(1) as cm1,
+    make_context_manager(2) as cm2,
+    make_context_manager(3) as cm3,
+    make_context_manager(4) as cm4,
+):
+    pass

tests/recorded/nested_with_nochange.txt

@@ -0,0 +1,75 @@
+with make_context_manager(1) as cm1:
+  pass
+
+# cannot mix `async with` and `with`
+async with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+      pass
+
+with make_context_manager(1) as cm1:
+    async with make_context_manager(2) as cm2:
+      pass
+
+# cannot have comments inside with statement
+with \
+     make_context_manager(1) as cm1, \
+     # comment in with statement
+     make_context_manager(2) as cm2 \
+:
+  with make_context_manager(3) as cm3:
+    pass
+
+with make_context_manager(1) as cm1:
+    # nested with
+    with make_context_manager(2) as cm2:
+        pass
+
+with make_context_manager(
+      1 # nested with
+    ) as cm1:
+    with make_context_manager(2) as cm2:
+        pass
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+        pass
+    a = "second statement blocks this refactor"
+
+================================================================================
+
+with make_context_manager(1) as cm1:
+  pass
+
+# cannot mix `async with` and `with`
+async with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+      pass
+
+with make_context_manager(1) as cm1:
+    async with make_context_manager(2) as cm2:
+      pass
+
+# cannot have comments inside with statement
+with \
+     make_context_manager(1) as cm1, \
+     # comment in with statement
+     make_context_manager(2) as cm2 \
+:
+  with make_context_manager(3) as cm3:
+    pass
+
+with make_context_manager(1) as cm1:
+    # nested with
+    with make_context_manager(2) as cm2:
+        pass
+
+with make_context_manager(
+      1 # nested with
+    ) as cm1:
+    with make_context_manager(2) as cm2:
+        pass
+
+with make_context_manager(1) as cm1:
+    with make_context_manager(2) as cm2:
+        pass
+    a = "second statement blocks this refactor"

tests/recorded/pybetter.txt

@@ -4,10 +4,6 @@ a == True
 a == False
 df[df.flag == True]  # doesn't break Pandas code
 
-with a:
-    with b:
-        pass
-
 with a:
     # comment
     with b:
@@ -21,10 +17,6 @@ a == True
 a == False
 df[df.flag == True]  # doesn't break Pandas code
 
-with a:
-    with b:
-        pass
-
 with a:
     # comment
     with b:

tests/test_expected_output.py

@@ -1,6 +1,7 @@
 """Update and check saved examples of shed formatting."""
 
 import pathlib
+import re
 import warnings
 
 import pytest
@@ -29,6 +30,10 @@ def test_saved_examples(filename: pathlib.Path, min_version):
     You can therefore see what changed by examining the `git diff`
     and roll back with `git reset`.
     """
+    stem = filename.stem
+    if re.search(r"_3\d+$", stem) and not stem.endswith(f"_3{min_version[1]}"):
+        pytest.skip(reason="Requires a different min-version spec.")
+
     joiner = "\n\n" + "=" * 80 + "\n\n"
     input_, expected, *_ = map(str.strip, (filename.read_text() + joiner).split(joiner))
     if filename.suffix == ".py" and "invalid" not in filename.stem: