fix: handle decorators after snakecode (#185)

snakemake · Mar 15, 2023 · 32d6c53 · 32d6c53
1 parent b5aa660
commit 32d6c53
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/snakefmt/parser/parser.py b/snakefmt/parser/parser.py
@@ -310,7 +310,7 @@ def get_next_queriable(self, snakefile: TokenIterator) -> Status:
                 else:
                     buffer += TAB * self.effective_indent
 
-            if token.type == tokenize.NAME and self.queriable:
+            if (token.type == tokenize.NAME or token.string == "@") and self.queriable:
                 self.queriable = False
                 return Status(
                     token, block_indent, self.cur_indent, buffer, False, pythonable

diff --git a/tests/test_formatter.py b/tests/test_formatter.py
@@ -377,6 +377,22 @@ def test_line_wrapped_python_code_inside_rule(self):
 
         assert actual == expected
 
+    def test_decorator_is_handled_correctly(self):
+        """https://github.com/snakemake/snakefmt/issues/144"""
+        snakecode = (
+            "# comment\n"
+            "rule foo:\n"
+            f"{TAB * 1}run:\n"
+            f'{TAB * 2}print("")\n\n\n'
+            "@contextlib.contextmanager\n"
+            "def f(wildcards):\n"
+            f"{TAB * 1}pass\n"
+        )
+        formatter = setup_formatter(snakecode)
+
+        actual = formatter.get_formatted()
+        assert actual == snakecode
+
 
 class TestComplexPythonFormatting:
     """