WIP use Flags to mark doctype tokens

Author: lagru

src/docstub/_doctype.py

@@ -1,5 +1,7 @@
 """Parsing of doctypes"""
 
+import enum
+import itertools
 import logging
 from collections.abc import Iterable
 from dataclasses import dataclass
@@ -35,31 +37,61 @@ def insert_between(iterable, *, sep):
     return out[:-1]
 
 
+class TokenFlag(enum.Flag):
+    # docstub: off
+    NAME = enum.auto()
+    NATLANG = enum.auto()
+    SUBSCRIPT = enum.auto()
+    LITERAL = enum.auto()
+    GENERATOR = enum.auto()
+    ARRAY = enum.auto()
+    UNION = enum.auto()
+    START = enum.auto()
+    STOP = enum.auto()
+    SEP = enum.auto()
+    # docstub: on
+
+    @classmethod
+    def _missing_(cls, value):
+        forbidden = {
+            *itertools.combinations([cls.START, cls.STOP, cls.SEP, cls.NAME], 2)
+        }
+        for pair in forbidden:
+            if value is (pair[0].value | pair[1].value):
+                raise ValueError(f"{pair[0].name}|{pair[1].name} not allowed")
+        return super()._missing_(value)
+
+
 class Token(str):
     """A token representing an atomic part of a doctype."""
 
-    __slots__ = ("value", "kind")
+    flag = TokenFlag
+
+    __slots__ = ("value", "kind", "pos")
 
-    def __new__(cls, value, *, kind):
+    def __new__(cls, value, *, kind, pos=None):
         self = super().__new__(cls, value)
-        self.kind = kind
+        self.kind = TokenFlag(kind)
+        self.pos = pos
         return self
 
     def __repr__(self):
         return f"{type(self).__name__}('{self}', kind={self.kind!r})"
 
     @classmethod
     def find_iter(cls, iterable, *, kind):
+        kind = TokenFlag(kind)
         for item in flatten_recursive(iterable):
-            if isinstance(item, cls) and item.kind == kind:
+            if isinstance(item, cls) and all(k & item.kind for k in kind):
                 yield item
 
     @classmethod
     def find_one(cls, iterable, *, kind):
         matching = list(cls.find_iter(iterable, kind=kind))
         if len(matching) != 1:
             msg = (
-                f"expected exactly one {cls.__name__} with {kind=}, got {len(matching)}"
+                f"expected exactly one {cls.__name__} with {kind=}, "
+                f"got {len(matching)}: {matching}"
             )
             raise ValueError(msg)
         return matching[0]
@@ -79,7 +111,11 @@ def qualname(self, tree):
         """
         children = tree.children
         _qualname = ".".join(children)
-        _qualname = Token(_qualname, kind="qualname")
+        _qualname = Token(
+            _qualname,
+            kind=Token.flag.NAME,
+            pos=(tree.meta.start_pos, tree.meta.end_pos),
+        )
         return _qualname
 
     def rst_role(self, tree):
@@ -92,7 +128,7 @@ def rst_role(self, tree):
         -------
         out : lark.Token
         """
-        qualname = Token.find_one(tree.children, kind="qualname")
+        qualname = Token.find_one(tree.children, kind=Token.flag.NAME)
         return qualname
 
     def union(self, tree):
@@ -105,7 +141,7 @@ def union(self, tree):
         -------
         out : list[str]
         """
-        sep = Token(" | ", kind="union_sep")
+        sep = Token(" | ", kind=Token.flag.UNION | Token.flag.SEP)
         out = insert_between(tree.children, sep=sep)
         return out
 
@@ -119,7 +155,7 @@ def subscription(self, tree):
         -------
         out : str
         """
-        return self._format_subscription(tree.children, name="subscription")
+        return self._format_subscription(tree.children)
 
     def natlang_literal(self, tree):
         """
@@ -131,8 +167,13 @@ def natlang_literal(self, tree):
         -------
         out : str
         """
-        items = [Token("Literal", kind="qualname"), *tree.children]
-        out = self._format_subscription(items, "nl_literal")
+        items = [
+            Token("Literal", kind=Token.flag.LITERAL | Token.flag.NAME),
+            *tree.children,
+        ]
+        out = self._format_subscription(
+            items, kind=Token.flag.LITERAL | Token.flag.NATLANG
+        )
 
         if len(tree.children) == 1:
             logger.warning(
@@ -143,6 +184,14 @@ def natlang_literal(self, tree):
             )
         return out
 
+    def literal_item(self, tree):
+        item, *other = tree.children
+        assert not other
+        kind = Token.flag.LITERAL
+        if isinstance(item, Token):
+            kind |= item.kind
+        return Token(item, kind=kind, pos=(tree.meta.start_pos, tree.meta.end_pos))
+
     def natlang_container(self, tree):
         """
         Parameters
@@ -153,7 +202,7 @@ def natlang_container(self, tree):
         -------
         out : str
         """
-        return self._format_subscription(tree.children, name="nl_container")
+        return self._format_subscription(tree.children, kind=Token.flag.NATLANG)
 
     def natlang_array(self, tree):
         """
@@ -165,11 +214,15 @@ def natlang_array(self, tree):
         -------
         out : str
         """
-        array_name = Token.find_one(tree.children, kind="array_name")
+        array_name = Token.find_one(
+            tree.children, kind=Token.flag.ARRAY | Token.flag.NAME
+        )
         items = tree.children.copy()
         items.remove(array_name)
-        items.insert(0, Token(array_name, kind="qualname"))
-        return self._format_subscription(items, name="nl_array")
+        items.insert(0, array_name)
+        return self._format_subscription(
+            items, kind=Token.flag.ARRAY | Token.flag.NATLANG
+        )
 
     def array_name(self, tree):
         """
@@ -186,7 +239,7 @@ def array_name(self, tree):
         # This currently relies on a hack that only allows specific names
         # in `array_expression` (see `ARRAY_NAME` terminal in gramar)
         qualname = self.qualname(tree)
-        qualname = Token(qualname, kind="array_name")
+        qualname = Token(qualname, kind=Token.flag.NAME | Token.flag.ARRAY)
         return qualname
 
     def shape(self, tree):
@@ -228,22 +281,24 @@ def extra_info(self, tree):
         logger.debug("dropping extra info")
         return lark.Discard
 
-    def _format_subscription(self, sequence, name):
-        sep = Token(", ", kind=f"{name}_sep")
+    def _format_subscription(self, sequence, kind=None):
+        if kind is None:
+            kind = Token.flag.SUBSCRIPT
+        else:
+            kind |= Token.flag.SUBSCRIPT
+
+        sep = Token(", ", kind=kind | Token.flag.SEP)
         container, *content = sequence
         content = insert_between(content, sep=sep)
         assert content
         out = [
             container,
-            Token("[", kind=f"{name}_start"),
+            Token("[", kind=kind | Token.flag.START),
             *content,
-            Token("]", kind=f"{name}_stop"),
+            Token("]", kind=kind | Token.flag.STOP),
         ]
         return out
 
-    def __default_token__(self, token):
-        return Token(token.value, kind=token.type.lower())
-
 
 @dataclass(frozen=True, slots=True)
 class ParsedDoctype:
@@ -265,16 +320,38 @@ def parse(cls, doctype):
 
         Examples
         --------
-        >>> ParsedDoctype.parse("tuple of int or ndarray of dtype (float or int)")
+        >>> doctype = ParsedDoctype.parse(
+        ...     "tuple of int or ndarray of dtype (float or int)"
+        ... )
+        >>> doctype
         <ParsedDoctype: 'tuple[int] | ndarray[float | int]'>
+        >>> doctype.qualnames
+        (Token('tuple', kind='qualname'),
+         Token('int', kind='qualname'),
+         Token('ndarray', kind='qualname'),
+         Token('float', kind='qualname'),
+         Token('int', kind='qualname'))
         """
         tree = _lark.parse(doctype)
-        result = DoctypeTransformer().transform(tree=tree)
-        result = tuple(flatten_recursive(result))
-        return cls(result, raw_doctype=doctype)
+        tokens = DoctypeTransformer().transform(tree=tree)
+        tokens = tuple(flatten_recursive(tokens))
+        return cls(tokens, raw_doctype=doctype)
 
     def __str__(self):
         return "".join(self.tokens)
 
     def __repr__(self):
-        return f"<{type(self).__name__}: '{self}'>"
+        return f"<{type(self).__name__} '{self}'>"
+
+    @property
+    def qualnames(self):
+        return tuple(Token.find_iter(self.tokens, kind=Token.flag.NAME))
+
+    def print_map_tokens_to_raw(self):
+        for token in self.tokens:
+            if token.pos is not None:
+                start, stop = token.pos
+                print(self.raw_doctype)
+                print(" " * start + "^" * (stop - start))
+                print(" " * start + token)
+                print()

src/docstub/doctype.lark

@@ -58,7 +58,7 @@ natlang_literal: "{" literal_item ("," literal_item)* "}"
 
 // An single item in a literal expression (or `optional`). We must also allow
 // for qualified names, since a "class" or enum can be used as a literal too.
-?literal_item: ELLIPSES | STRING | SIGNED_NUMBER | qualname
+literal_item: ELLIPSES | STRING | SIGNED_NUMBER | qualname
 
 
 // Natural language forms of the subscription expression for containers.