From 95b218fbede4bf933c2c6a3c6f99dee07256b1de Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Fri, 29 May 2026 18:45:06 +0200
Subject: [PATCH] tests: add tests for python parser

---
 tests/cases/python-parser/01_simple_types.py  |  14 ++
 .../python-parser/01_simple_types.py.ref.json |  85 +++++++++
 tests/cases/python-parser/02_custom_types.py  |  29 ++++
 .../python-parser/02_custom_types.py.ref.json | 162 ++++++++++++++++++
 tests/cases/python-parser/03_functions.py     |  15 ++
 .../python-parser/03_functions.py.ref.json    | 149 ++++++++++++++++
 6 files changed, 454 insertions(+)
 create mode 100644 tests/cases/python-parser/01_simple_types.py
 create mode 100644 tests/cases/python-parser/01_simple_types.py.ref.json
 create mode 100644 tests/cases/python-parser/02_custom_types.py
 create mode 100644 tests/cases/python-parser/02_custom_types.py.ref.json
 create mode 100644 tests/cases/python-parser/03_functions.py
 create mode 100644 tests/cases/python-parser/03_functions.py.ref.json

diff --git a/tests/cases/python-parser/01_simple_types.py b/tests/cases/python-parser/01_simple_types.py
new file mode 100644
index 0000000..3566f9a
--- /dev/null
+++ b/tests/cases/python-parser/01_simple_types.py
@@ -0,0 +1,14 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+df: Frame[
+    verified: bool,
+    birth_year: int,
+    height: float + ( _ > 0 ) + ( _ < 250 ),
+    name: str,
+    date: datetime,
+    float,
+    unknown: _,
+    _
+]
diff --git a/tests/cases/python-parser/01_simple_types.py.ref.json b/tests/cases/python-parser/01_simple_types.py.ref.json
new file mode 100644
index 0000000..e4fd591
--- /dev/null
+++ b/tests/cases/python-parser/01_simple_types.py.ref.json
@@ -0,0 +1,85 @@
+{
+  "stmts": [
+    {
+      "_type": "TypeAssign",
+      "name": "df",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "verified",
+            "type": {
+              "_type": "BaseType",
+              "base": "bool",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "birth_year",
+            "type": {
+              "_type": "BaseType",
+              "base": "int",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "height",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "(_ > 0) + (_ < 250)"
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "name",
+            "type": {
+              "_type": "BaseType",
+              "base": "str",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "date",
+            "type": {
+              "_type": "BaseType",
+              "base": "datetime",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": null,
+            "type": {
+              "_type": "BaseType",
+              "base": "float",
+              "param": null
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "unknown",
+            "type": null
+          },
+          {
+            "_type": "FrameColumn",
+            "name": null,
+            "type": {
+              "_type": "BaseType",
+              "base": "_",
+              "param": null
+            }
+          }
+        ]
+      }
+    }
+  ]
+}
\ No newline at end of file
diff --git a/tests/cases/python-parser/02_custom_types.py b/tests/cases/python-parser/02_custom_types.py
new file mode 100644
index 0000000..32dbab3
--- /dev/null
+++ b/tests/cases/python-parser/02_custom_types.py
@@ -0,0 +1,29 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+import midas
+
+midas.using("02_custom_types.midas")
+
+df: Frame[
+    location: GeoLocation
+]
+
+lat: Column[GeoLocation] = df["location"].lat
+lon: Column[GeoLocation] = df["location"].lon
+
+lat + lon
+
+lat1: Latitude = lat[0]
+lat2: Latitude = lat[1]
+lat_diff: Difference[Latitude] = lat2 - lat1
+
+df2: Frame[
+    age: int + (_ >= 0),
+    height: float + (_ >= 0),
+]
+df2_bis: Frame[
+    age: int + Positive,
+    height: float + Positive,
+]
diff --git a/tests/cases/python-parser/02_custom_types.py.ref.json b/tests/cases/python-parser/02_custom_types.py.ref.json
new file mode 100644
index 0000000..49b861f
--- /dev/null
+++ b/tests/cases/python-parser/02_custom_types.py.ref.json
@@ -0,0 +1,162 @@
+{
+  "stmts": [
+    {
+      "_type": "ExpressionStmt",
+      "expr": {
+        "_type": "CallExpr",
+        "callee": {
+          "_type": "GetExpr",
+          "object": {
+            "_type": "VariableExpr",
+            "name": "midas"
+          },
+          "name": "using"
+        },
+        "arguments": [
+          {
+            "_type": "LiteralExpr",
+            "value": "02_custom_types.midas"
+          }
+        ],
+        "keywords": {}
+      }
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "df",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "location",
+            "type": {
+              "_type": "BaseType",
+              "base": "GeoLocation",
+              "param": null
+            }
+          }
+        ]
+      }
+    },
+    {
+      "_type": "ExpressionStmt",
+      "expr": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "lat"
+        },
+        "operator": "+",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "lon"
+        }
+      }
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "lat_diff",
+      "type": {
+        "_type": "BaseType",
+        "base": "Difference",
+        "param": {
+          "_type": "BaseType",
+          "base": "Latitude",
+          "param": null
+        }
+      }
+    },
+    {
+      "_type": "AssignStmt",
+      "targets": [
+        {
+          "_type": "VariableExpr",
+          "name": "lat_diff"
+        }
+      ],
+      "value": {
+        "_type": "BinaryExpr",
+        "left": {
+          "_type": "VariableExpr",
+          "name": "lat2"
+        },
+        "operator": "-",
+        "right": {
+          "_type": "VariableExpr",
+          "name": "lat1"
+        }
+      }
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "df2",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "age",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "int",
+                "param": null
+              },
+              "constraint": "_ >= 0"
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "height",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "_ >= 0"
+            }
+          }
+        ]
+      }
+    },
+    {
+      "_type": "TypeAssign",
+      "name": "df2_bis",
+      "type": {
+        "_type": "FrameType",
+        "columns": [
+          {
+            "_type": "FrameColumn",
+            "name": "age",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "int",
+                "param": null
+              },
+              "constraint": "Positive"
+            }
+          },
+          {
+            "_type": "FrameColumn",
+            "name": "height",
+            "type": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "Positive"
+            }
+          }
+        ]
+      }
+    }
+  ]
+}
\ No newline at end of file
diff --git a/tests/cases/python-parser/03_functions.py b/tests/cases/python-parser/03_functions.py
new file mode 100644
index 0000000..3b07899
--- /dev/null
+++ b/tests/cases/python-parser/03_functions.py
@@ -0,0 +1,15 @@
+# type: ignore
+# ruff: disable[F821]
+from __future__ import annotations
+
+
+def func(
+    col1: Column[float + (0 <= _ <= 1)],
+    col2: Column[float + (0 <= _ <= 1)],
+) -> Column[float + (0 <= _ <= 2)]:
+    result: Column[float + (0 <= _ <= 2)] = col1 + col2
+    return result
+
+
+def func2(a: int, /, b: float, *, c: str):
+    pass
diff --git a/tests/cases/python-parser/03_functions.py.ref.json b/tests/cases/python-parser/03_functions.py.ref.json
new file mode 100644
index 0000000..529455b
--- /dev/null
+++ b/tests/cases/python-parser/03_functions.py.ref.json
@@ -0,0 +1,149 @@
+{
+  "stmts": [
+    {
+      "_type": "Function",
+      "name": "func",
+      "posonlyargs": [],
+      "args": [
+        {
+          "name": "col1",
+          "type": {
+            "_type": "BaseType",
+            "base": "Column",
+            "param": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "0 <= _ <= 1"
+            }
+          },
+          "default": null
+        },
+        {
+          "name": "col2",
+          "type": {
+            "_type": "BaseType",
+            "base": "Column",
+            "param": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "0 <= _ <= 1"
+            }
+          },
+          "default": null
+        }
+      ],
+      "sink": null,
+      "kwonlyargs": [],
+      "kw_sink": null,
+      "returns": {
+        "_type": "BaseType",
+        "base": "Column",
+        "param": {
+          "_type": "ConstraintType",
+          "type": {
+            "_type": "BaseType",
+            "base": "float",
+            "param": null
+          },
+          "constraint": "0 <= _ <= 2"
+        }
+      },
+      "body": [
+        {
+          "_type": "TypeAssign",
+          "name": "result",
+          "type": {
+            "_type": "BaseType",
+            "base": "Column",
+            "param": {
+              "_type": "ConstraintType",
+              "type": {
+                "_type": "BaseType",
+                "base": "float",
+                "param": null
+              },
+              "constraint": "0 <= _ <= 2"
+            }
+          }
+        },
+        {
+          "_type": "AssignStmt",
+          "targets": [
+            {
+              "_type": "VariableExpr",
+              "name": "result"
+            }
+          ],
+          "value": {
+            "_type": "BinaryExpr",
+            "left": {
+              "_type": "VariableExpr",
+              "name": "col1"
+            },
+            "operator": "+",
+            "right": {
+              "_type": "VariableExpr",
+              "name": "col2"
+            }
+          }
+        },
+        {
+          "_type": "ReturnStmt",
+          "value": {
+            "_type": "VariableExpr",
+            "name": "result"
+          }
+        }
+      ]
+    },
+    {
+      "_type": "Function",
+      "name": "func2",
+      "posonlyargs": [
+        {
+          "name": "a",
+          "type": {
+            "_type": "BaseType",
+            "base": "int",
+            "param": null
+          },
+          "default": null
+        }
+      ],
+      "args": [
+        {
+          "name": "b",
+          "type": {
+            "_type": "BaseType",
+            "base": "float",
+            "param": null
+          },
+          "default": null
+        }
+      ],
+      "sink": null,
+      "kwonlyargs": [
+        {
+          "name": "c",
+          "type": {
+            "_type": "BaseType",
+            "base": "str",
+            "param": null
+          },
+          "default": null
+        }
+      ],
+      "kw_sink": null,
+      "returns": null,
+      "body": []
+    }
+  ]
+}
\ No newline at end of file