Skip to content

Commit

Permalink
Fixes #490: Ensure custom types support pickling.
Browse files Browse the repository at this point in the history
Previously `NamedRowTuple` did not roundtrip successfully.
Added tests for all custom types.
  • Loading branch information
ragnard committed Oct 15, 2024
1 parent caff0e8 commit 836de89
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 0 deletions.
35 changes: 35 additions & 0 deletions tests/unit/test_types.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pickle
from datetime import datetime, time
from decimal import Decimal

from trino import types

import pytest

IDENTITY = lambda v: v

type_instances = [
(types.Time(time(11, 47, 23), Decimal(0.314)), lambda v: v.to_python_type()),
(types.TimeWithTimeZone(time(11, 47, 23), Decimal(0.314)), lambda v: v.to_python_type()),
(types.Timestamp(datetime(2024, 10, 15, 11, 47, 23), Decimal(0.314)), lambda v: v.to_python_type()),
(types.TimestampWithTimeZone(datetime(2024, 10, 15, 11, 47, 23), Decimal(0.314)), lambda v: v.to_python_type()),
(types.NamedRowTuple(["Alice", 38], ["name", "age"], ["varchar", "integer"]), IDENTITY),
]

@pytest.mark.parametrize("value,fn", type_instances)
def test_pickle_roundtripping(value, fn):
bytes = pickle.dumps(value)
unpickled_value = pickle.loads(bytes)
assert fn(value) == fn(unpickled_value)
9 changes: 9 additions & 0 deletions trino/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,5 +111,14 @@ def __getattr__(self, name: str) -> Any:
if self._names.count(name):
raise ValueError("Ambiguous row field reference: " + name)

def __getnewargs__(self):
return (tuple(self), (), ())

def __getstate__(self):
return vars(self)

def __setstate__(self, state):
vars(self).update(state)

def __repr__(self) -> str:
return self._repr

0 comments on commit 836de89

Please sign in to comment.