Merge pull request #8658 from Tishj/pytype_optional

[Python] Support `Optional[...]` in DuckDBPyType
duckdb · Nov 8, 2023 · 85eed5c · 85eed5c
2 parents b0dbd9b + e199473
commit 85eed5c
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 3 deletions.
diff --git a/tools/pythonpkg/src/typing/pytype.cpp b/tools/pythonpkg/src/typing/pytype.cpp
@@ -199,10 +199,22 @@ static bool IsMapType(const py::tuple &args) {
 	return true;
 }
 
-static LogicalType FromUnionType(const py::object &obj) {
+static py::tuple FilterNones(const py::tuple &args) {
+	py::list result;
+
+	for (const auto &arg : args) {
+		py::object object = py::reinterpret_borrow<py::object>(arg);
+		if (object.is(py::none().get_type())) {
+			continue;
+		}
+		result.append(object);
+	}
+	return py::tuple(result);
+}
+
+static LogicalType FromUnionTypeInternal(const py::tuple &args) {
 	idx_t index = 1;
 	child_list_t<LogicalType> members;
-	py::tuple args = obj.attr("__args__");
 
 	for (const auto &arg : args) {
 		auto name = StringUtil::Format("u%d", index++);
@@ -211,6 +223,19 @@ static LogicalType FromUnionType(const py::object &obj) {
 	}
 
 	return LogicalType::UNION(std::move(members));
+}
+
+static LogicalType FromUnionType(const py::object &obj) {
+	py::tuple args = obj.attr("__args__");
+
+	// Optional inserts NoneType into the Union
+	// all types are nullable in DuckDB so we just filter the Nones
+	auto filtered_args = FilterNones(args);
+	if (filtered_args.size() == 1) {
+		// If only a single type is left, dont construct a UNION
+		return FromObject(filtered_args[0]);
+	}
+	return FromUnionTypeInternal(filtered_args);
 };
 
 static LogicalType FromGenericAlias(const py::object &obj) {

diff --git a/tools/pythonpkg/tests/fast/test_type.py b/tools/pythonpkg/tests/fast/test_type.py
@@ -2,7 +2,8 @@
 import os
 import pandas as pd
 import pytest
-from typing import Union
+from typing import Union, Optional
+import sys
 
 from duckdb.typing import (
     SQLNULL,
@@ -32,6 +33,7 @@
     BIT,
     INTERVAL,
 )
+import duckdb.typing
 
 
 class TestType(object):
@@ -187,3 +189,26 @@ def test_attribute_accessor(self):
 
         child_type = type.v2.child
         assert str(child_type) == 'MAP(BLOB, BIT)'
+
+    def test_optional(self):
+        type = duckdb.typing.DuckDBPyType(Optional[str])
+        assert type == 'VARCHAR'
+        type = duckdb.typing.DuckDBPyType(Optional[Union[int, bool]])
+        assert type == 'UNION(u1 BIGINT, u2 BOOLEAN)'
+        type = duckdb.typing.DuckDBPyType(Optional[list[int]])
+        assert type == 'BIGINT[]'
+        type = duckdb.typing.DuckDBPyType(Optional[dict[int, str]])
+        assert type == 'MAP(BIGINT, VARCHAR)'
+        type = duckdb.typing.DuckDBPyType(Optional[dict[Optional[int], Optional[str]]])
+        assert type == 'MAP(BIGINT, VARCHAR)'
+        type = duckdb.typing.DuckDBPyType(Optional[dict[Optional[int], Optional[str]]])
+        assert type == 'MAP(BIGINT, VARCHAR)'
+        type = duckdb.typing.DuckDBPyType(Optional[Union[Optional[str], Optional[bool]]])
+        assert type == 'UNION(u1 VARCHAR, u2 BOOLEAN)'
+        type = duckdb.typing.DuckDBPyType(Union[str, None])
+        assert type == 'VARCHAR'
+
+    @pytest.mark.skipif(sys.version_info < (3, 10), reason="'str | None' syntax requires Python 3.10 or higher")
+    def test_optional_310(self):
+        type = duckdb.typing.DuckDBPyType(str | None)
+        assert type == 'VARCHAR'