From 69941e482e2b5e7d6e2a2a2c4713a74767aa5216 Mon Sep 17 00:00:00 2001 From: Hendrik Makait Date: Fri, 11 Aug 2023 15:00:00 +0200 Subject: [PATCH 1/3] Make ToPickle generic and add typing --- distributed/protocol/serialize.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py index b2ca44b4752..0493702865b 100644 --- a/distributed/protocol/serialize.py +++ b/distributed/protocol/serialize.py @@ -7,7 +7,7 @@ from enum import Enum from functools import partial from types import ModuleType -from typing import Any, Literal +from typing import Any, Generic, Literal, TypeVar import msgpack @@ -28,6 +28,8 @@ ) from distributed.utils import ensure_memoryview, has_keyword +T = TypeVar("T") + dask_serialize = dask.utils.Dispatch("dask_serialize") dask_deserialize = dask.utils.Dispatch("dask_deserialize") @@ -561,7 +563,7 @@ def __ne__(self, other): return not (self == other) -class ToPickle: +class ToPickle(Generic[T]): """Mark an object that should be pickled Both the scheduler and workers with automatically unpickle this @@ -572,19 +574,21 @@ class ToPickle: to False, the scheduler will raise an exception instead. """ - def __init__(self, data): + data: T + + def __init__(self, data: T) -> None: self.data = data - def __repr__(self): + def __repr__(self) -> str: return "" % str(self.data) - def __eq__(self, other): + def __eq__(self, other: object) -> bool: return isinstance(other, type(self)) and other.data == self.data - def __ne__(self, other): + def __ne__(self, other: object) -> bool: return not (self == other) - def __hash__(self): + def __hash__(self) -> int: return hash(self.data) From 791913c0b22dd2c1a4c08c546507ce5e29b65a03 Mon Sep 17 00:00:00 2001 From: Hendrik Makait Date: Mon, 14 Aug 2023 16:15:56 +0200 Subject: [PATCH 2/3] Update distributed/protocol/serialize.py Co-authored-by: crusaderky --- distributed/protocol/serialize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py index 0493702865b..97b6b04902f 100644 --- a/distributed/protocol/serialize.py +++ b/distributed/protocol/serialize.py @@ -576,7 +576,7 @@ class ToPickle(Generic[T]): data: T - def __init__(self, data: T) -> None: + def __init__(self, data: T): self.data = data def __repr__(self) -> str: From 0d10798d0fc5e1e5b617be147cecd193a3266734 Mon Sep 17 00:00:00 2001 From: Hendrik Makait Date: Mon, 14 Aug 2023 16:17:14 +0200 Subject: [PATCH 3/3] Remove __ne__ --- distributed/protocol/serialize.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/distributed/protocol/serialize.py b/distributed/protocol/serialize.py index 97b6b04902f..42503001dca 100644 --- a/distributed/protocol/serialize.py +++ b/distributed/protocol/serialize.py @@ -585,9 +585,6 @@ def __repr__(self) -> str: def __eq__(self, other: object) -> bool: return isinstance(other, type(self)) and other.data == self.data - def __ne__(self, other: object) -> bool: - return not (self == other) - def __hash__(self) -> int: return hash(self.data)