From b08bdab75f6a4c37992c02e6e29bbbf0af9352a0 Mon Sep 17 00:00:00 2001
From: Darren Burns <darrenb900@gmail.com>
Date: Mon, 7 Feb 2022 15:18:07 +0000
Subject: [PATCH] Change whitespace trimming during css var sub, add more tests

---
 src/textual/css/parse.py | 33 +++++++++++++++++++++------------
 tests/css/test_parse.py  | 24 ++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/src/textual/css/parse.py b/src/textual/css/parse.py
index 955f19762..af600d262 100644
--- a/src/textual/css/parse.py
+++ b/src/textual/css/parse.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from collections import defaultdict
 from functools import lru_cache
-from typing import Iterator, Iterable
+from typing import Iterator, Iterable, Optional
 
 from rich import print
 from rich.cells import cell_len
@@ -204,10 +204,6 @@ def parse_declarations(css: str, path: str) -> Styles:
     return styles_builder.styles
 
 
-def _is_whitespace(token: Token) -> bool:
-    return token.name == "whitespace"
-
-
 def _unresolved(
     variable_name: str, location: tuple[int, int]
 ) -> UnresolvedVariableError:
@@ -241,16 +237,21 @@ def substitute_references(tokens: Iterator[Token]) -> Iterable[Token]:
             variable_name = token.value[1:-1]  # Trim the $ and the :, i.e. "$x:" -> "x"
             yield token
 
-            # Store the tokens for any variable definitions, and substitute
-            # any variable references we encounter with them.
-            leading_whitespace = True
             while True:
                 token = next(tokens, None)
+                if token.name == "whitespace":
+                    yield token
+                else:
+                    break
+
+            # Store the tokens for any variable definitions, and substitute
+            # any variable references we encounter with them.
+            while True:
                 if not token:
                     break
-                elif leading_whitespace and token.name == "whitespace":
+                elif token.name == "whitespace":
+                    variables[variable_name].append(token)
                     yield token
-                    leading_whitespace = False
                 elif token.name == "variable_value_end":
                     yield token
                     break
@@ -263,8 +264,8 @@ def substitute_references(tokens: Iterator[Token]) -> Iterable[Token]:
                         variable_tokens.extend(reference_tokens)
                         ref_location = token.location
                         ref_length = cell_len(token.value)
-                        for token in reference_tokens:
-                            yield token.with_reference(
+                        for _token in reference_tokens:
+                            yield _token.with_reference(
                                 ReferencedBy(
                                     name=ref_name,
                                     location=ref_location,
@@ -278,6 +279,7 @@ def substitute_references(tokens: Iterator[Token]) -> Iterable[Token]:
                 else:
                     variables[variable_name].append(token)
                     yield token
+                token = next(tokens, None)
         elif token.name == "variable_ref":
             variable_name = token.value[1:]  # Trim the $, so $x -> x
             if variable_name in variables:
@@ -299,6 +301,13 @@ def substitute_references(tokens: Iterator[Token]) -> Iterable[Token]:
 
 
 def parse(css: str, path: str) -> Iterable[RuleSet]:
+    """Parse CSS by tokenizing it, performing variable substitution,
+    and generating rule sets from it.
+
+    Args:
+        css (str): The input CSS
+        path (str): Path to the CSS
+    """
     tokens = iter(substitute_references(tokenize(css, path)))
     while True:
         token = next(tokens, None)
diff --git a/tests/css/test_parse.py b/tests/css/test_parse.py
index 490f8f7fa..5142d5cc9 100644
--- a/tests/css/test_parse.py
+++ b/tests/css/test_parse.py
@@ -148,6 +148,30 @@ class TestVariableReferenceSubstitution:
             Token(name='declaration_set_end', value='}', path='', code=css, location=(1, 24), referenced_by=None)
         ]
 
+    def test_variable_definition_eof(self):
+        css = "$x: 1"
+        assert list(substitute_references(tokenize(css, ""))) == [
+            Token(name='variable_name', value='$x:', path='', code=css, location=(0, 0), referenced_by=None),
+            Token(name='whitespace', value=' ', path='', code=css, location=(0, 3), referenced_by=None),
+            Token(name='number', value='1', path='', code=css, location=(0, 4), referenced_by=None)
+        ]
+
+    def test_variable_reference_whitespace_trimming(self):
+        css = "$x:    123;.thing{border: $x}"
+        assert list(substitute_references(tokenize(css, ""))) == [
+            Token(name='variable_name', value='$x:', path='', code=css, location=(0, 0), referenced_by=None),
+            Token(name='whitespace', value='    ', path='', code=css, location=(0, 3), referenced_by=None),
+            Token(name='number', value='123', path='', code=css, location=(0, 7), referenced_by=None),
+            Token(name='variable_value_end', value=';', path='', code=css, location=(0, 10), referenced_by=None),
+            Token(name='selector_start_class', value='.thing', path='', code=css, location=(0, 11), referenced_by=None),
+            Token(name='declaration_set_start', value='{', path='', code=css, location=(0, 17), referenced_by=None),
+            Token(name='declaration_name', value='border:', path='', code=css, location=(0, 18), referenced_by=None),
+            Token(name='whitespace', value=' ', path='', code=css, location=(0, 25), referenced_by=None),
+            Token(name='number', value='123', path='', code=css, location=(0, 7),
+                  referenced_by=ReferencedBy(name='x', location=(0, 26), length=2)),
+            Token(name='declaration_set_end', value='}', path='', code=css, location=(0, 28), referenced_by=None)
+        ]
+
 
 class TestParseLayout:
     def test_valid_layout_name(self):