python-trio · Zac-HD · May 1, 2023 · Apr 30, 2023
diff --git a/flake8_trio/visitors/helpers.py b/flake8_trio/visitors/helpers.py
@@ -377,7 +377,11 @@ def flatten_preserving_comments(node: cst.BaseCompoundStatement):
     # and IndentedBlock
     if isinstance(node.body, cst.SimpleStatementSuite):
         # `with ...: pass;pass;pass` -> pass;pass;pass
-        return cst.SimpleStatementLine(node.body.body, leading_lines=new_leading_lines)
+        return cst.SimpleStatementLine(
+            node.body.body,
+            leading_lines=new_leading_lines,
+            trailing_whitespace=node.body.trailing_whitespace,
+        )
 
     assert isinstance(node.body, cst.IndentedBlock)
     nodes = list(node.body.body)

diff --git a/tests/autofix_files/trio100.py b/tests/autofix_files/trio100.py
@@ -9,7 +9,7 @@
 
 async def function_name():
     # fmt: off
-    ...; ...; ...
+    ...; ...; ...  # error: 15, "trio", "fail_after"
     # fmt: on
     # error: 15, "trio", "fail_after"
     ...

diff --git a/tests/autofix_files/trio100.py.diff b/tests/autofix_files/trio100.py.diff
@@ -13,7 +13,7 @@
  async def function_name():
      # fmt: off
 -    async with trio.fail_after(10): ...; ...; ...  # error: 15, "trio", "fail_after"
-+    ...; ...; ...
++    ...; ...; ...  # error: 15, "trio", "fail_after"
      # fmt: on
 -    async with trio.fail_after(10):  # error: 15, "trio", "fail_after"
 -        ...

diff --git a/tests/autofix_files/trio100_simple_autofix.py b/tests/autofix_files/trio100_simple_autofix.py
@@ -25,7 +25,7 @@
 # a
 # b
 # fmt: off
-...;...;...
+...;...;... # error: 5, "trio", "move_on_after"
 # fmt: on
 # c
 # d
@@ -51,3 +51,9 @@
 # c
 ...; ...; ...
 # fmt: on
+
+
+# same-line with
+# fmt: off
+print(1)  # TRIO100: 5, 'trio', 'fail_after'
+# fmt: on
diff --git a/tests/autofix_files/trio100_simple_autofix.py.diff b/tests/autofix_files/trio100_simple_autofix.py.diff
@@ -1,6 +1,6 @@
 ---
 +++
-@@ x,50 x,51 @@
+@@ x,56 x,57 @@
 
  # a
  # b
@@ -41,7 +41,7 @@
  # b
  # fmt: off
 -with trio.move_on_after(10): ...;...;... # error: 5, "trio", "move_on_after"
-+...;...;...
++...;...;... # error: 5, "trio", "move_on_after"
  # fmt: on
  # c
  # d
@@ -83,3 +83,10 @@
 +# c
 +...; ...; ...
  # fmt: on
+
+
+ # same-line with
+ # fmt: off
+-with trio.fail_after(5): print(1)  # TRIO100: 5, 'trio', 'fail_after'
++print(1)  # TRIO100: 5, 'trio', 'fail_after'
+ # fmt: on
diff --git a/tests/eval_files/trio100_simple_autofix.py b/tests/eval_files/trio100_simple_autofix.py
@@ -50,3 +50,9 @@
     # c
 ): ...; ...; ...
 # fmt: on
+
+
+# same-line with
+# fmt: off
+with trio.fail_after(5): print(1)  # TRIO100: 5, 'trio', 'fail_after'
+# fmt: on
diff --git a/tests/test_flake8_trio.py b/tests/test_flake8_trio.py
@@ -164,7 +164,11 @@ def check_autofix(
         " update the autofix files."
     )
     # and assert that the diff is the same, which it should be if the above passes
-    assert added_autofix_diff == autofix_diff_content
+    assert added_autofix_diff == autofix_diff_content, (
+        "THIS SHOULD NOT HAPPEN: diff in the autofix diff - without there being a diff"
+        " in the autofixed code, run with --generate-autofix if the test file has"
+        " changed to update the autofix files."
+    )
 
 
 # This can be further cleaned up by adding the other return values from