Fix #93 (2)

Treat special edge case where comments follow snakemake directives
nested in python code.

snakefmt/formatter.py

@@ -10,7 +10,7 @@
 from snakefmt.exceptions import InvalidParameterSyntax, InvalidPython
 from snakefmt.logging import Warnings
 from snakefmt.parser.grammar import SnakeRule
-from snakefmt.parser.parser import Parser
+from snakefmt.parser.parser import Parser, comment_start
 from snakefmt.parser.syntax import (
     COMMENT_SPACING,
     TAB,
@@ -32,10 +32,6 @@
 )
 
 
-def comment_start(string: str) -> bool:
-    return string.lstrip().startswith("#")
-
-
 class Formatter(Parser):
     def __init__(
         self,

snakefmt/parser/parser.py

@@ -33,6 +33,10 @@ def denext(self, token: Token) -> None:
         self._buffered_tokens.append(token)
 
 
+def comment_start(string: str) -> bool:
+    return string.lstrip().startswith("#")
+
+
 class Parser(ABC):
     def __init__(self, snakefile: TokenIterator):
         self.grammar = Grammar(
@@ -71,14 +75,14 @@ def __init__(self, snakefile: TokenIterator):
                 self.context.code_indent = None
                 status = self.process_keyword(status, self.from_python)
             else:
-                if not self.context.accepts_python_code and not keyword[0] == "#":
+                if not self.context.accepts_python_code and not comment_start(keyword):
                     raise SyntaxError(
                         f"L{status.token.start[0]}: Unrecognised keyword '{keyword}' "
                         f"in {self.context.keyword_name} definition"
                     )
                 else:
                     self.buffer += keyword
-                    if self.context.code_indent is None:
+                    if self.context.code_indent is None and not comment_start(keyword):
                         self.context.code_indent = status.indent
                     status = self.context.get_next_queriable(self.snakefile)
                     self.buffer += status.buffer

tests/test_formatter.py

@@ -389,6 +389,17 @@ def test_spaced_out_consecutive_dedented_directive(self):
         formatter = setup_formatter(snakecode)
         assert formatter.get_formatted() == snakecode
 
+    def test_comment_support_after_python_code(self):
+        snakecode = (
+            'if config["a"]:\n\n'
+            f'{TAB * 1}include: "module_a.smk"\n\n\n'
+            f'# include: "module_b.smk"\n\n\n'
+            f'if config["c"]:\n\n'
+            f'{TAB * 1}include: "module_c.smk"\n'
+        )
+        formatter = setup_formatter(snakecode)
+        assert formatter.get_formatted() == snakecode
+
     def test_nested_ifelse_statements(self):
         snakecode = (
             'if config["a"] is None:\n\n'