Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Python] Support Optional[...] in DuckDBPyType #8658

Merged
merged 5 commits into from Nov 8, 2023
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
29 changes: 27 additions & 2 deletions tools/pythonpkg/src/typing/pytype.cpp
Expand Up @@ -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++);
Expand All @@ -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) {
Expand Down
19 changes: 18 additions & 1 deletion tools/pythonpkg/tests/fast/test_type.py
Expand Up @@ -2,7 +2,7 @@
import os
import pandas as pd
import pytest
from typing import Union
from typing import Union, Optional

from duckdb.typing import (
SQLNULL,
Expand Down Expand Up @@ -32,6 +32,7 @@
BIT,
INTERVAL,
)
import duckdb.typing


class TestType(object):
Expand Down Expand Up @@ -186,3 +187,19 @@ def test_attribute_accessor(self):

child_type = type.v2.child
assert str(child_type) == 'MAP(BLOB, BIT)'

def test_optional(self):
Tishj marked this conversation as resolved.
Show resolved Hide resolved
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)'