From 1423e6b35377ef2431aa7400155ac69b10afb12e Mon Sep 17 00:00:00 2001
From: Brice Letcher <bletcher@ebi.ac.uk>
Date: Sat, 10 Dec 2022 19:07:20 +0100
Subject: [PATCH] Fix #159

---
 snakefmt/formatter.py     | 11 +++++++----
 snakefmt/parser/parser.py |  3 +--
 tests/test_formatter.py   |  2 +-
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/snakefmt/formatter.py b/snakefmt/formatter.py
index 50a3988..1a56d49 100644
--- a/snakefmt/formatter.py
+++ b/snakefmt/formatter.py
@@ -30,6 +30,7 @@
     r"(.*)^(if|elif|else|with|for|while)([^:]*)(:.*)", re.S | re.M
 )
 after_if_keywords = ("elif", "else")
+is_all_comments = lambda string: all(map(comment_start, string.splitlines()))
 
 
 class Formatter(Parser):
@@ -105,8 +106,7 @@ def flush_buffer(
             if comment_start(self.buffer.rstrip().splitlines()[-1]):
                 formatted += "\n"
         # Only stick together separated single-parm keywords when separated by comments
-        buffer_is_all_comments = all(map(comment_start, self.buffer.splitlines()))
-        if not buffer_is_all_comments:
+        if not is_all_comments(self.buffer):
             self.last_recognised_keyword = ""
         self.add_newlines(self.block_indent, formatted, final_flush, in_global_context)
         self.buffer = ""
@@ -135,8 +135,11 @@ def process_keyword_param(
     def run_black_format_str(
         self, string: str, target_indent: int, extra_spacing: int = 0
     ) -> str:
-        needs_artifical_nest = self.from_python and not string.startswith("if")
-        needs_artifical_nest = False
+        needs_artifical_nest = (
+            self.from_python
+            and not string.startswith("if")
+            and not is_all_comments(string)
+        )
         if needs_artifical_nest:
             string = f"if x:\n{textwrap.indent(string, TAB)}"
 
diff --git a/snakefmt/parser/parser.py b/snakefmt/parser/parser.py
index 3397ec6..363677e 100644
--- a/snakefmt/parser/parser.py
+++ b/snakefmt/parser/parser.py
@@ -88,8 +88,6 @@ def __init__(self, snakefile: TokenIterator):
                     in_if_else = self.buffer.startswith(("if", "else", "elif"))
                     if self.syntax.from_python or status.pythonable or in_if_else:
                         self.from_python = True
-                    else:  # Over-indented context gets reset
-                        self.syntax.cur_indent = max(self.keyword_indent - 1, 0)
                 elif self.from_python:
                     # We are exiting python context, so force spacing out keywords
                     self.last_recognised_keyword = ""
@@ -117,6 +115,7 @@ def __init__(self, snakefile: TokenIterator):
                         self.from_python
                         and status.cur_indent == 0
                         and not self.last_block_was_snakecode
+                        and self.block_indent > 0
                     ):
                         # This flushes any nested python code following a
                         # nested snakemake keyword
diff --git a/tests/test_formatter.py b/tests/test_formatter.py
index 0aec4ba..d6aaa78 100644
--- a/tests/test_formatter.py
+++ b/tests/test_formatter.py
@@ -961,7 +961,7 @@ def test_snakecode_after_indented_comment_does_not_get_unindented(self):
         snakecode = (
             'if config.get("s3_dst"):\n\n'
             f'{TAB * 1}include: "workflow/rule1.smk"\n'
-            f'{TAB * 1}include: "workflow/rule2.smk"\n\n'
+            f'{TAB * 1}include: "workflow/rule2.smk"\n'
             f"{TAB * 1}# a comment\n"
             f"{TAB * 1}# further comment\n"
             f'{TAB * 1}include: "workflow/rule3.smk"\n'