Merge pull request #19 from microsoft/tokenization_fixes

niclasko · web-flow · commit 4d2cc6d6fca8 · 2026-02-07T13:03:36.000+01:00
Tokenization fixes
diff --git a/flowquery-py/src/parsing/parser.py b/flowquery-py/src/parsing/parser.py
@@ -326,7 +326,7 @@ def _parse_create(self) -> Optional[Operation]:
             if not self.token.is_colon():
                 raise ValueError("Expected ':' for relationship type")
             self.set_next_token()
-            if not self.token.is_identifier():
+            if not self.token.is_identifier_or_keyword():
                 raise ValueError("Expected relationship type identifier")
             rel_type = self.token.value or ""
             self.set_next_token()
@@ -450,17 +450,17 @@ def _parse_node(self) -> Optional[Node]:
         self.set_next_token()
         self._skip_whitespace_and_comments()
         identifier: Optional[str] = None
-        if self.token.is_identifier():
+        if self.token.is_identifier_or_keyword():
             identifier = self.token.value
             self.set_next_token()
         self._skip_whitespace_and_comments()
         label: Optional[str] = None
         peek = self.peek()
-        if not self.token.is_colon() and peek is not None and peek.is_identifier():
+        if not self.token.is_colon() and peek is not None and peek.is_identifier_or_keyword():
             raise ValueError("Expected ':' for node label")
-        if self.token.is_colon() and (peek is None or not peek.is_identifier()):
+        if self.token.is_colon() and (peek is None or not peek.is_identifier_or_keyword()):
             raise ValueError("Expected node label identifier")
-        if self.token.is_colon() and peek is not None and peek.is_identifier():
+        if self.token.is_colon() and peek is not None and peek.is_identifier_or_keyword():
             self.set_next_token()
             label = cast(str, self.token.value)  # Guaranteed by is_identifier check
             self.set_next_token()
@@ -495,13 +495,13 @@ def _parse_relationship(self) -> Optional[Relationship]:
             return None
         self.set_next_token()
         variable: Optional[str] = None
-        if self.token.is_identifier():
+        if self.token.is_identifier_or_keyword():
             variable = self.token.value
             self.set_next_token()
         if not self.token.is_colon():
             raise ValueError("Expected ':' for relationship type")
         self.set_next_token()
-        if not self.token.is_identifier():
+        if not self.token.is_identifier_or_keyword():
             raise ValueError("Expected relationship type identifier")
         rel_type: str = self.token.value or ""
         self.set_next_token()
@@ -633,14 +633,14 @@ def _parse_expressions(
     def _parse_operand(self, expression: Expression) -> bool:
         """Parse a single operand (without operators). Returns True if an operand was parsed."""
         self._skip_whitespace_and_comments()
-        if self.token.is_identifier() and (self.peek() is None or not self.peek().is_left_parenthesis()):
+        if self.token.is_identifier_or_keyword() and (self.peek() is None or not self.peek().is_left_parenthesis()):
             identifier = self.token.value or ""
             reference = Reference(identifier, self._variables.get(identifier))
             self.set_next_token()
             lookup = self._parse_lookup(reference)
             expression.add_node(lookup)
             return True
-        elif self.token.is_identifier() and self.peek() is not None and self.peek().is_left_parenthesis():
+        elif self.token.is_identifier_or_keyword() and self.peek() is not None and self.peek().is_left_parenthesis():
             func = self._parse_predicate_function() or self._parse_function()
             if func is not None:
                 lookup = self._parse_lookup(func)
@@ -650,7 +650,7 @@ def _parse_operand(self, expression: Expression) -> bool:
             self.token.is_left_parenthesis()
             and self.peek() is not None
             and (
-                self.peek().is_identifier()
+                self.peek().is_identifier_or_keyword()
                 or self.peek().is_colon()
                 or self.peek().is_right_parenthesis()
             )
@@ -734,7 +734,7 @@ def _parse_lookup(self, node: ASTNode) -> ASTNode:
         while True:
             if self.token.is_dot():
                 self.set_next_token()
-                if not self.token.is_identifier() and not self.token.is_keyword():
+                if not self.token.is_identifier_or_keyword():
                     raise ValueError("Expected identifier")
                 lookup = Lookup()
                 lookup.index = Identifier(self.token.value or "")
@@ -847,7 +847,7 @@ def _parse_alias(self) -> Optional[Alias]:
         self._expect_previous_token_to_be_whitespace_or_comment()
         self.set_next_token()
         self._expect_and_skip_whitespace_and_comments()
-        if not self.token.is_identifier():
+        if not self.token.is_identifier_or_keyword():
             raise ValueError("Expected identifier")
         alias = Alias(self.token.value or "")
         self.set_next_token()
diff --git a/flowquery-py/src/tokenization/token.py b/flowquery-py/src/tokenization/token.py
@@ -106,6 +106,24 @@ def IDENTIFIER(value: str) -> Token:
     def is_identifier(self) -> bool:
         return self._type == TokenType.IDENTIFIER or self._type == TokenType.BACKTICK_STRING
 
+    def is_keyword_that_cannot_be_identifier(self) -> bool:
+        """Returns True for keywords that have special expression-level roles
+        and should not be treated as identifiers (NULL, CASE, WHEN, THEN, ELSE, END)."""
+        return self.is_keyword() and (
+            self.is_null()
+            or self.is_case()
+            or self.is_when()
+            or self.is_then()
+            or self.is_else()
+            or self.is_end()
+        )
+
+    def is_identifier_or_keyword(self) -> bool:
+        """Returns True if the token is an identifier or a keyword that can be used as an identifier."""
+        return self.is_identifier() or (
+            self.is_keyword() and not self.is_keyword_that_cannot_be_identifier()
+        )
+
     # String token
 
     @staticmethod
diff --git a/flowquery-py/tests/compute/test_runner.py b/flowquery-py/tests/compute/test_runner.py
@@ -1584,4 +1584,63 @@ async def test_schema_returns_nodes_and_relationships_with_sample_data(self):
         assert chases["sample"] is not None
         assert "left_id" not in chases["sample"]
         assert "right_id" not in chases["sample"]
-        assert "speed" in chases["sample"]
+        assert "speed" in chases["sample"]
+
+    @pytest.mark.asyncio
+    async def test_reserved_keywords_as_identifiers(self):
+        """Test reserved keywords as identifiers."""
+        runner = Runner("""
+            WITH 1 AS return
+            RETURN return
+        """)
+        await runner.run()
+        results = runner.results
+        assert len(results) == 1
+        assert results[0]["return"] == 1
+
+    @pytest.mark.asyncio
+    async def test_reserved_keywords_as_parts_of_identifiers(self):
+        """Test reserved keywords as parts of identifiers."""
+        runner = Runner("""
+            unwind [
+                {from: "Alice", to: "Bob", organizer: "Charlie"},
+                {from: "Bob", to: "Charlie", organizer: "Alice"},
+                {from: "Charlie", to: "Alice", organizer: "Bob"}
+            ] as data
+            return data.from as from, data.to as to, data.organizer as organizer
+        """)
+        await runner.run()
+        results = runner.results
+        assert len(results) == 3
+        assert results[0] == {"from": "Alice", "to": "Bob", "organizer": "Charlie"}
+        assert results[1] == {"from": "Bob", "to": "Charlie", "organizer": "Alice"}
+        assert results[2] == {"from": "Charlie", "to": "Alice", "organizer": "Bob"}
+
+    @pytest.mark.asyncio
+    async def test_reserved_keywords_as_relationship_types_and_labels(self):
+        """Test reserved keywords as relationship types and labels."""
+        await Runner("""
+            CREATE VIRTUAL (:Return) AS {
+                unwind [
+                    {id: 1, name: 'Node 1'},
+                    {id: 2, name: 'Node 2'}
+                ] as record
+                RETURN record.id as id, record.name as name
+            }
+        """).run()
+        await Runner("""
+            CREATE VIRTUAL (:Return)-[:With]-(:Return) AS {
+                unwind [
+                    {left_id: 1, right_id: 2}
+                ] as record
+                RETURN record.left_id as left_id, record.right_id as right_id
+            }
+        """).run()
+        runner = Runner("""
+            MATCH (a:Return)-[:With]->(b:Return)
+            RETURN a.name AS name1, b.name AS name2
+        """)
+        await runner.run()
+        results = runner.results
+        assert len(results) == 1
+        assert results[0] == {"name1": "Node 1", "name2": "Node 2"}
diff --git a/flowquery-py/tests/parsing/test_parser.py b/flowquery-py/tests/parsing/test_parser.py
@@ -719,3 +719,12 @@ def test_relationship_with_properties(self):
         assert isinstance(relationship, Relationship)
         assert relationship.properties.get("since") is not None
         assert relationship.properties["since"].value() == 2022
+
+    def test_case_statement_with_keywords_as_identifiers(self):
+        """Test that CASE/WHEN/THEN/ELSE/END are not treated as identifiers."""
+        parser = Parser()
+        ast = parser.parse("RETURN CASE WHEN 1 THEN 2 ELSE 3 END")
+        assert "Case" in ast.print()
+        assert "When" in ast.print()
+        assert "Then" in ast.print()
+        assert "Else" in ast.print()
diff --git a/flowquery-py/tests/tokenization/test_tokenizer.py b/flowquery-py/tests/tokenization/test_tokenizer.py
@@ -162,3 +162,37 @@ def test_relationship_with_hops(self):
         tokens = tokenizer.tokenize()
         assert tokens is not None
         assert len(tokens) > 0
+
+    def test_reserved_keywords_as_identifiers(self):
+        """Test reserved keywords as identifiers."""
+        tokenizer = Tokenizer("""
+            WITH 1 AS return
+            RETURN return
+        """)
+        tokens = tokenizer.tokenize()
+        assert tokens is not None
+        assert len(tokens) > 0
+
+    def test_reserved_keywords_as_part_of_identifiers(self):
+        """Test reserved keywords as part of identifiers."""
+        tokenizer = Tokenizer("""
+            unwind [
+                {from: "Alice", to: "Bob", organizer: "Charlie"},
+                {from: "Bob", to: "Charlie", organizer: "Alice"},
+                {from: "Charlie", to: "Alice", organizer: "Bob"}
+            ] as data
+            return data.from, data.to
+        """)
+        tokens = tokenizer.tokenize()
+        assert tokens is not None
+        assert len(tokens) > 0
+
+    def test_reserved_keywords_as_relationship_types_and_labels(self):
+        """Test reserved keywords as relationship types and labels."""
+        tokenizer = Tokenizer("""
+            MATCH (a:RETURN)-[r:WITH]->(b:RETURN)
+            RETURN a, b
+        """)
+        tokens = tokenizer.tokenize()
+        assert tokens is not None
+        assert len(tokens) > 0
diff --git a/src/parsing/parser.ts b/src/parsing/parser.ts
@@ -353,7 +353,7 @@ class Parser extends BaseParser {
                 throw new Error("Expected ':' for relationship type");
             }
             this.setNextToken();
-            if (!this.token.isIdentifier()) {
+            if (!this.token.isIdentifierOrKeyword()) {
                 throw new Error("Expected relationship type identifier");
             }
             const type: string = this.token.value || "";
@@ -412,19 +412,19 @@ class Parser extends BaseParser {
         this.setNextToken();
         this.skipWhitespaceAndComments();
         let identifier: string | null = null;
-        if (this.token.isIdentifier()) {
+        if (this.token.isIdentifierOrKeyword()) {
             identifier = this.token.value || "";
             this.setNextToken();
         }
         this.skipWhitespaceAndComments();
         let label: string | null = null;
-        if (!this.token.isColon() && this.peek()?.isIdentifier()) {
+        if (!this.token.isColon() && this.peek()?.isIdentifierOrKeyword()) {
             throw new Error("Expected ':' for node label");
         }
-        if (this.token.isColon() && !this.peek()?.isIdentifier()) {
+        if (this.token.isColon() && !this.peek()?.isIdentifierOrKeyword()) {
             throw new Error("Expected node label identifier");
         }
-        if (this.token.isColon() && this.peek()?.isIdentifier()) {
+        if (this.token.isColon() && this.peek()?.isIdentifierOrKeyword()) {
             this.setNextToken();
             label = this.token.value || "";
             this.setNextToken();
@@ -595,15 +595,15 @@ class Parser extends BaseParser {
         }
         this.setNextToken();
         let variable: string | null = null;
-        if (this.token.isIdentifier()) {
+        if (this.token.isIdentifierOrKeyword()) {
             variable = this.token.value || "";
             this.setNextToken();
         }
         if (!this.token.isColon()) {
             throw new Error("Expected ':' for relationship type");
         }
         this.setNextToken();
-        if (!this.token.isIdentifier()) {
+        if (!this.token.isIdentifierOrKeyword()) {
             throw new Error("Expected relationship type identifier");
         }
         const type: string = this.token.value || "";
@@ -745,14 +745,14 @@ class Parser extends BaseParser {
      */
     private parseOperand(expression: Expression): boolean {
         this.skipWhitespaceAndComments();
-        if (this.token.isIdentifier() && !this.peek()?.isLeftParenthesis()) {
+        if (this.token.isIdentifierOrKeyword() && !this.peek()?.isLeftParenthesis()) {
             const identifier: string = this.token.value || "";
             const reference = new Reference(identifier, this.variables.get(identifier));
             this.setNextToken();
             const lookup = this.parseLookup(reference);
             expression.addNode(lookup);
             return true;
-        } else if (this.token.isIdentifier() && this.peek()?.isLeftParenthesis()) {
+        } else if (this.token.isIdentifierOrKeyword() && this.peek()?.isLeftParenthesis()) {
             const func = this.parsePredicateFunction() || this.parseFunction();
             if (func !== null) {
                 const lookup = this.parseLookup(func);
@@ -761,7 +761,7 @@ class Parser extends BaseParser {
             }
         } else if (
             this.token.isLeftParenthesis() &&
-            (this.peek()?.isIdentifier() ||
+            (this.peek()?.isIdentifierOrKeyword() ||
                 this.peek()?.isColon() ||
                 this.peek()?.isRightParenthesis())
         ) {
@@ -857,7 +857,7 @@ class Parser extends BaseParser {
         while (true) {
             if (this.token.isDot()) {
                 this.setNextToken();
-                if (!this.token.isIdentifier() && !this.token.isKeyword()) {
+                if (!this.token.isIdentifierOrKeyword()) {
                     throw new Error("Expected identifier");
                 }
                 lookup = new Lookup();
diff --git a/src/tokenization/token.ts b/src/tokenization/token.ts
@@ -104,6 +104,10 @@ class Token {
         return this._type === TokenType.IDENTIFIER || this._type === TokenType.BACKTICK_STRING;
     }
 
+    public isIdentifierOrKeyword(): boolean {
+        return this.isIdentifier() || (this.isKeyword() && !this.isKeywordThatCannotBeIdentifier());
+    }
+
     // String token
 
     public static STRING(value: string, quoteChar: string = '"'): Token {
@@ -387,6 +391,18 @@ class Token {
         return this._type === TokenType.KEYWORD;
     }
 
+    public isKeywordThatCannotBeIdentifier(): boolean {
+        return (
+            this.isKeyword() &&
+            (this.isNull() ||
+                this.isCase() ||
+                this.isWhen() ||
+                this.isThen() ||
+                this.isElse() ||
+                this.isEnd())
+        );
+    }
+
     public static get WITH(): Token {
         return new Token(TokenType.KEYWORD, Keyword.WITH);
     }
diff --git a/tests/compute/runner.test.ts b/tests/compute/runner.test.ts
diff --git a/tests/parsing/parser.test.ts b/tests/parsing/parser.test.ts
diff --git a/tests/tokenization/tokenizer.test.ts b/tests/tokenization/tokenizer.test.ts
