From bdaae76310bda0edc71fe156d8d550c319ddf7b1 Mon Sep 17 00:00:00 2001
From: "Oriol (ProDesk)" <oriol.abril.pla@gmail.com>
Date: Sat, 22 Jun 2024 19:46:22 +0200
Subject: [PATCH] draft for pandas examples support

---
 examples/example_pkg-stubs/_basic.pyi | 11 ++++++++++-
 examples/example_pkg/_basic.py        | 25 ++++++++++++++++++++++++-
 src/docstub/doctype.lark              |  9 ++++++---
 3 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/examples/example_pkg-stubs/_basic.pyi b/examples/example_pkg-stubs/_basic.pyi
index 1d7e2db..a7fceff 100644
--- a/examples/example_pkg-stubs/_basic.pyi
+++ b/examples/example_pkg-stubs/_basic.pyi
@@ -14,10 +14,19 @@ def func_empty(a1: Any, a2: Any, a3: Any) -> None: ...
 def func_contains(
     self,
     a1: list[float],
-    a2: dict[str, Union[int, str]],
+    a2: sequence[int] | float,
     a3: Sequence[int | float],
     a4: frozenset[bytes],
 ) -> tuple[tuple[int, ...], list[int]]: ...
+def func_contains_dict(
+    self,
+    a1: dict[["str", "int | str"]],
+    a2: dict[str, Union[int, str]],
+    a3: mapping[["int", "str"]],
+) -> dict[["int | str", "float"]]: ...
+def func_literals(
+    self, a1: Literal["A", "B", "C"], a2: Literal[0, "index", 1, "columns", None]
+) -> None: ...
 def func_literals(
     a1: Literal[1, 3, "foo"], a2: Literal["uno", 2, "drei", "four"] = ...
 ) -> None: ...
diff --git a/examples/example_pkg/_basic.py b/examples/example_pkg/_basic.py
index 3933b94..f62fbf0 100644
--- a/examples/example_pkg/_basic.py
+++ b/examples/example_pkg/_basic.py
@@ -32,7 +32,7 @@ def func_contains(self, a1, a2, a3, a4):
     Parameters
     ----------
     a1 : list[float]
-    a2 : dict[str, Union[int, str]]
+    a2 : sequence of int or float
     a3 : Sequence[int | float]
     a4 : frozenset[bytes]
 
@@ -42,6 +42,29 @@ def func_contains(self, a1, a2, a3, a4):
     r2 : list of int
     """
 
+def func_contains_dict(self, a1, a2, a3):
+    """Dummy.
+
+    Parameters
+    ----------
+    a1 : dict of {str : int or str}
+    a2 : dict[str, Union[int, str]]
+    a3 : mapping of {int : str}
+
+    Returns
+    -------
+    r1 : dict of {int or str : float}
+    """
+
+def func_literals(self, a1, a2):
+    """Dummy.
+
+    Parameters
+    ----------
+    a1 : {"A", "B", "C"}
+    a2 : {0 or "index", 1 or "columns", None}, default None
+    """
+
 
 def func_literals(a1, a2="uno"):
     """Dummy
diff --git a/src/docstub/doctype.lark b/src/docstub/doctype.lark
index 7cc4f95..e0fb16f 100644
--- a/src/docstub/doctype.lark
+++ b/src/docstub/doctype.lark
@@ -1,11 +1,12 @@
 ?start : doctype
 
-doctype : type_or ("," optional)? ("," extra_info)?
+doctype : (literals | type_or) ("," optional)? ("," extra_info)?
 
 type_or : type (("or" | "|") type)*
 
+literals : "{" literal (("," | "or") literal)* "}"
+
 ?type : qualname
-      | "{" literal ("," literal)* "}" -> literals
       | container_of
       | shape_n_dtype
 
@@ -23,7 +24,9 @@ contains: "[" type_or ("," type_or)* "]"
 
 
 // Container-of
-container_of : NAME "of" type_or
+container_of : NAME "of" ( type_or | dict_subtypes )
+
+dict_subtypes : "{" type_or ":" type_or "}"
 
 
 // Array-like form with dtype or shape information