Skip to content

Commit 8de3d07

Browse files
committed
fix the 'description' attribute, using DuckDBPyType
1 parent 3a9043d commit 8de3d07

File tree

4 files changed

+53
-4
lines changed

4 files changed

+53
-4
lines changed

duckdb/__init__.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,26 @@ def version():
1818
"functional"
1919
])
2020

21+
class DBAPITypeObject:
22+
def __init__(self, name, types):
23+
self.name = name
24+
self.types = list(types)
25+
26+
def __eq__(self, other):
27+
if isinstance(other, typing.DuckDBPyType):
28+
return other in self.types
29+
return False
30+
31+
def __repr__(self):
32+
return f"<DBAPITypeObject {self.name}>"
33+
34+
# Define the standard DBAPI sentinels
35+
STRING = DBAPITypeObject("STRING", {"VARCHAR", "CHAR", "TEXT"})
36+
NUMBER = DBAPITypeObject("NUMBER", {"INTEGER", "BIGINT", "DECIMAL", "DOUBLE"})
37+
DATETIME = DBAPITypeObject("DATETIME", {"DATE", "TIME", "TIMESTAMP"})
38+
BINARY = DBAPITypeObject("BINARY", {"BLOB"})
39+
ROWID = None
40+
2141
# Classes
2242
from _duckdb import (
2343
DuckDBPyRelation,

src/duckdb_py/pyresult.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -579,7 +579,7 @@ py::list DuckDBPyResult::GetDescription(const vector<string> &names, const vecto
579579

580580
for (idx_t col_idx = 0; col_idx < names.size(); col_idx++) {
581581
auto py_name = py::str(names[col_idx]);
582-
auto py_type = GetTypeToPython(types[col_idx]);
582+
auto py_type = DuckDBPyType(types[col_idx]);
583583
desc.append(py::make_tuple(py_name, py_type, py::none(), py::none(), py::none(), py::none(), py::none()));
584584
}
585585
return desc;

src/duckdb_py/typing/pytype.cpp

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -325,8 +325,17 @@ static LogicalType FromObject(const py::object &object) {
325325
void DuckDBPyType::Initialize(py::handle &m) {
326326
auto type_module = py::class_<DuckDBPyType, shared_ptr<DuckDBPyType>>(m, "DuckDBPyType", py::module_local());
327327

328-
type_module.def("__repr__", &DuckDBPyType::ToString, "Stringified representation of the type object");
329-
type_module.def("__eq__", &DuckDBPyType::Equals, "Compare two types for equality", py::arg("other"));
328+
type_module.def("__repr__", &DuckDBPyType::ToString, "Stringified representation of the type object")
329+
.def("__eq__", [](const DuckDBPyType &self, py::handle other) -> py::object {
330+
if (py::isinstance<DuckDBPyType>(other)) {
331+
return py::bool_(self.Equals(other.cast<shared_ptr<DuckDBPyType>>()));
332+
}
333+
else if (py::isinstance<py::str>(other)) {
334+
return py::bool_(self.EqualsString(other.cast<string>()));
335+
}
336+
// Return NotImplemented for unsupported types
337+
return py::reinterpret_borrow<py::object>(Py_NotImplemented);
338+
});
330339
type_module.def("__eq__", &DuckDBPyType::EqualsString, "Compare two types for equality", py::arg("other"));
331340
type_module.def_property_readonly("id", &DuckDBPyType::GetId);
332341
type_module.def_property_readonly("children", &DuckDBPyType::Children);

tests/fast/api/test_duckdb_connection.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import duckdb
2+
import duckdb.typing
23
import pytest
34
from conftest import NumpyPandas, ArrowPandas
45

@@ -113,9 +114,28 @@ def test_readonly_properties(self):
113114
duckdb.execute("select 42")
114115
description = duckdb.description()
115116
rowcount = duckdb.rowcount()
116-
assert description == [('42', 'NUMBER', None, None, None, None, None)]
117+
assert description == [('42', 'INTEGER', None, None, None, None, None)]
117118
assert rowcount == -1
118119

120+
def test_description(self):
121+
duckdb.execute("select 42 a, 'test' b, true c")
122+
types = [x[1] for x in duckdb.description()]
123+
124+
STRING = duckdb.STRING
125+
NUMBER = duckdb.NUMBER
126+
DATETIME = duckdb.DATETIME
127+
128+
assert(types[1] == STRING)
129+
assert(STRING == types[1])
130+
assert(types[0] != STRING)
131+
assert((types[1] != STRING) == False)
132+
assert((STRING != types[1]) == False)
133+
134+
assert(types[1] in [STRING])
135+
assert(types[1] in [STRING, NUMBER])
136+
assert(types[1] not in [NUMBER, DATETIME])
137+
138+
119139
def test_execute(self):
120140
assert [([4, 2],)] == duckdb.execute("select [4,2]").fetchall()
121141

0 commit comments

Comments
 (0)