From 8195ed035b0720218cdd5ac80904d9650fb9f8e0 Mon Sep 17 00:00:00 2001
From: Sam Dudley <dudley.co.uk@gmail.com>
Date: Fri, 29 Dec 2023 23:03:11 +0000
Subject: [PATCH] refactor compiler to fix failing formatting tests

---
 justfile                    |  4 +--
 python_html_dsl/compiler.py | 56 +++++++++++++++++++------------------
 python_html_dsl/tokens.py   |  9 ++++++
 tests/test_dsl.py           | 33 +++++++++++++++++++---
 4 files changed, 69 insertions(+), 33 deletions(-)

diff --git a/justfile b/justfile
index 0238625..aa15f52 100644
--- a/justfile
+++ b/justfile
@@ -16,8 +16,8 @@ check:
 	poetry run ruff .
 	poetry run mypy .
 
-test:
-	poetry run pytest --cov=python_html_dsl
+test *args:
+	poetry run pytest --cov=python_html_dsl {{args}}
 
 coverage:
 	poetry run coverage html
diff --git a/python_html_dsl/compiler.py b/python_html_dsl/compiler.py
index 80f6bf0..341895d 100644
--- a/python_html_dsl/compiler.py
+++ b/python_html_dsl/compiler.py
@@ -3,7 +3,6 @@
 from typing import TYPE_CHECKING, Any
 
 from .tokens import ClosingTag, Content, OpeningTag, Token
-from .utils import is_block_tag, is_self_closing_tag
 
 if TYPE_CHECKING:
     from .types import HtmlAttributes
@@ -14,12 +13,12 @@ def compile(self, tokens: deque[Token]) -> str:
         self.tokens = tokens
 
         self.depth: int = 0
-        self.stack: deque[str] = deque()
         self.token: Token | None = tokens.popleft()
         self.code: deque[str] = deque()
 
         while self.token:
             getattr(self, f"visit_{self.token.__class__.__name__}")(self.token)
+            self.move_cursor(self.token, self.peek())
             self.eat()
 
         if self.code[-1] != "\n":
@@ -33,6 +32,12 @@ def eat(self) -> None:
         except IndexError:
             self.token = None
 
+    def peek(self) -> Token | None:
+        try:
+            return self.tokens[0]
+        except IndexError:
+            return None
+
     def append(self, fragment: str) -> None:
         self.code.append(fragment)
 
@@ -42,40 +47,37 @@ def indent(self) -> None:
     def newline(self) -> None:
         self.append("\n")
 
-    def in_block(self) -> bool:
-        try:
-            return is_block_tag(self.stack[-1])
-        except IndexError:
-            return False
-
     def visit_OpeningTag(self, tag: OpeningTag) -> None:
-        if self.in_block():
-            self.newline()
-            self.indent()
-
         self.append(f"<{tag.name}{self.render_attrs(tag.attrs)}>")
-        self.stack.append(tag.name)
-
-        if is_block_tag(tag.name) and not is_self_closing_tag(tag.name):
-            self.depth += 1
 
     def visit_Content(self, content: Content) -> None:
-        if self.in_block():
-            self.newline()
-            self.indent()
-
         text = escape(content.text) if not content.safe else content.text
-
         self.append(text)
 
     def visit_ClosingTag(self, tag: ClosingTag) -> None:
-        if is_block_tag(tag.name):
-            self.depth -= 1
-            self.newline()
-            self.indent()
-
         self.append(f"</{tag.name}>")
-        self.stack.pop()
+
+    def move_cursor(self, token: Token, next_token: Token | None) -> None:
+        if not next_token:
+            return
+
+        match (token, next_token):
+            case [OpeningTag(), ClosingTag()]:
+                self.depth = self.depth
+            case [OpeningTag(), _]:
+                self.depth += 1
+            case [_, ClosingTag()]:
+                self.depth -= 1
+
+        match (token, next_token):
+            # current or next token is an opening block tag
+            case [OpeningTag(is_block=True), _] | [_, OpeningTag(is_block=True)]:
+                self.newline()
+                self.indent()
+            # current or next token is a closing block tag
+            case [ClosingTag(is_block=True), _] | [_, ClosingTag(is_block=True)]:
+                self.newline()
+                self.indent()
 
     @classmethod
     def render_attrs(cls, attrs: "HtmlAttributes") -> str:
diff --git a/python_html_dsl/tokens.py b/python_html_dsl/tokens.py
index 5bc05a5..3b31dcd 100644
--- a/python_html_dsl/tokens.py
+++ b/python_html_dsl/tokens.py
@@ -1,4 +1,5 @@
 from .types import HtmlAttributes
+from .utils import is_block_tag
 
 
 class Token:
@@ -11,6 +12,10 @@ def __init__(self, name: str, attrs: HtmlAttributes):
         self.name = name
         self.attrs = attrs
 
+    @property
+    def is_block(self) -> bool:
+        return is_block_tag(self.name)
+
     def __repr__(self) -> str:
         return f"OpeningTag({self.name!r}, {self.attrs!r})"
 
@@ -28,5 +33,9 @@ class ClosingTag(Token):
     def __init__(self, name: str):
         self.name = name
 
+    @property
+    def is_block(self) -> bool:
+        return is_block_tag(self.name)
+
     def __repr__(self) -> str:
         return f"ClosingTag({self.name!r})"
diff --git a/tests/test_dsl.py b/tests/test_dsl.py
index 9356160..884211e 100644
--- a/tests/test_dsl.py
+++ b/tests/test_dsl.py
@@ -129,9 +129,11 @@ def test_nested_inline_block(self) -> None:
         html = render(h("span", [h("div", ["hello"])]))
         assert html == dedent(
             """\
-            <span><div>
-                hello
-            </div></span>
+            <span>
+                <div>
+                    hello
+                </div>
+            </span>
             """
         )
 
@@ -180,7 +182,17 @@ def test_nested_inline_flat_div(self) -> None:
             """
         )
 
-    def test_inline_then_text(self) -> None:
+    def test_block_inline_inline(self) -> None:
+        html = h("div", [h("span", "foo"), h("span", "bar")])
+        assert render(html) == dedent(
+            """\
+            <div>
+                <span>foo</span><span>bar</span>
+            </div>
+            """
+        )
+
+    def test_block_inline_text(self) -> None:
         html = h("div", [h("span", "foo"), "bar"])
         assert render(html) == dedent(
             """\
@@ -189,3 +201,16 @@ def test_inline_then_text(self) -> None:
             </div>
             """
         )
+
+    def test_block_following_content(self) -> None:
+        html = h("div", ["hello", h("div", "world")])
+        assert render(html) == dedent(
+            """\
+            <div>
+                hello
+                <div>
+                    world
+                </div>
+            </div>
+            """
+        )