Skip to content

Commit

Permalink
fix compile and import
Browse files Browse the repository at this point in the history
  • Loading branch information
@chaokunyang
chaokunyang committed Dec 17, 2024
1 parent d2fb76e commit 473c796
Show file tree
Hide file tree
Showing 7 changed files with 144 additions and 145 deletions.
16 changes: 6 additions & 10 deletions python/pyfury/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,20 +27,16 @@
except ImportError:
ENABLE_FURY_CYTHON_SERIALIZATION = False

from pyfury._fury import ( # noqa: F401,F403,F811 # pylint: disable=unused-import
Fury,
Language,
OpaqueObject,
)
if ENABLE_FURY_CYTHON_SERIALIZATION:
from pyfury._serialization import ( # noqa: F401,F811
Fury,
Language,
ClassInfo,
OpaqueObject,
)
else:
from pyfury._fury import ( # noqa: F401,F403,F811 # pylint: disable=unused-import
Fury,
Language,
ClassInfo,
OpaqueObject,
)

from pyfury._struct import ( # noqa: F401,F403,F811 # pylint: disable=unused-import
ComplexObjectSerializer,
)
Expand Down
60 changes: 49 additions & 11 deletions python/pyfury/_fury.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,16 +23,6 @@
from dataclasses import dataclass
from typing import Union, Iterable

from pyfury._serializer import SerializationContext
from pyfury._registry import (
PYINT_CLASS_ID,
PYBOOL_CLASS_ID,
STRING_CLASS_ID,
NOT_NULL_STRING_FLAG,
NOT_NULL_PYINT_FLAG,
NOT_NULL_PYBOOL_FLAG,
NO_CLASS_ID,
)
from pyfury.buffer import Buffer
from pyfury.resolver import (
MapRefResolver,
Expand All @@ -54,8 +44,26 @@
logger = logging.getLogger(__name__)


DEFAULT_DYNAMIC_WRITE_STRING_ID = -1
MAGIC_NUMBER = 0x62D4
DEFAULT_DYNAMIC_WRITE_STRING_ID = -1
DYNAMIC_TYPE_ID = -1
USE_CLASSNAME = 0
USE_CLASS_ID = 1
# preserve 0 as flag for class id not set in ClassInfo`
NO_CLASS_ID = 0
PYINT_CLASS_ID = 1
PYFLOAT_CLASS_ID = 2
PYBOOL_CLASS_ID = 3
STRING_CLASS_ID = 4
PICKLE_CLASS_ID = 5
PICKLE_STRONG_CACHE_CLASS_ID = 6
PICKLE_CACHE_CLASS_ID = 7
# `NOT_NULL_VALUE_FLAG` + `CLASS_ID << 1` in little-endian order
NOT_NULL_PYINT_FLAG = NOT_NULL_VALUE_FLAG & 0b11111111 | (PYINT_CLASS_ID << 9)
NOT_NULL_PYFLOAT_FLAG = NOT_NULL_VALUE_FLAG & 0b11111111 | (PYFLOAT_CLASS_ID << 9)
NOT_NULL_PYBOOL_FLAG = NOT_NULL_VALUE_FLAG & 0b11111111 | (PYBOOL_CLASS_ID << 9)
NOT_NULL_STRING_FLAG = NOT_NULL_VALUE_FLAG & 0b11111111 | (STRING_CLASS_ID << 9)
SMALL_STRING_THRESHOLD = 16


class Language(enum.Enum):
Expand Down Expand Up @@ -475,6 +483,36 @@ def reset(self):
self.reset_read()


class SerializationContext:
"""
A context is used to add some context-related information, so that the
serializers can setup relation between serializing different objects.
The context will be reset after finished serializing/deserializing the
object tree.
"""

__slots__ = ("objects",)

def __init__(self):
self.objects = dict()

def add(self, key, obj):
self.objects[id(key)] = obj

def __contains__(self, key):
return id(key) in self.objects

def __getitem__(self, key):
return self.objects[id(key)]

def get(self, key):
return self.objects.get(id(key))

def reset(self):
if len(self.objects) > 0:
self.objects.clear()


_ENABLE_CLASS_REGISTRATION_FORCIBLY = os.getenv(
"ENABLE_CLASS_REGISTRATION_FORCIBLY", "0"
) in {
Expand Down
171 changes: 79 additions & 92 deletions python/pyfury/_registry.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,13 @@
from pyfury.lib import mmh3
from pyfury.serializer import (
Serializer,
EnumSerializer,
PickleSerializer,
Numpy1DArraySerializer,
NDArraySerializer,
PyArraySerializer,
DynamicPyArraySerializer,
NoneSerializer,
_PickleStub,
PickleStrongCacheStub,
PickleCacheStub,
PickleStrongCacheSerializer,
NoneSerializer,
BooleanSerializer,
ByteSerializer,
Expand All @@ -45,14 +41,13 @@
SetSerializer,
EnumSerializer,
SliceSerializer,
PickleCacheSerializer,
PickleStrongCacheSerializer,
PickleSerializer,
)
from pyfury._struct import ComplexObjectSerializer
from pyfury.buffer import Buffer
from pyfury.meta.metastring import MetaStringEncoder, MetaStringDecoder
from pyfury.resolver import (
NULL_FLAG,
NOT_NULL_VALUE_FLAG,
)
from pyfury.type import (
TypeId,
Int8Type,
Expand All @@ -63,37 +58,35 @@
Float64Type,
load_class,
)
from pyfury._fury import (
DEFAULT_DYNAMIC_WRITE_STRING_ID,
DYNAMIC_TYPE_ID,
USE_CLASSNAME,
USE_CLASS_ID,
# preserve 0 as flag for class id not set in ClassInfo`
NO_CLASS_ID,
PYINT_CLASS_ID,
PYFLOAT_CLASS_ID,
PYBOOL_CLASS_ID,
STRING_CLASS_ID,
PICKLE_CLASS_ID,
PICKLE_STRONG_CACHE_CLASS_ID,
PICKLE_CACHE_CLASS_ID,
NOT_NULL_PYINT_FLAG,
NOT_NULL_PYFLOAT_FLAG,
NOT_NULL_PYBOOL_FLAG,
NOT_NULL_STRING_FLAG,
SMALL_STRING_THRESHOLD,
)

try:
import numpy as np
except ImportError:
np = None


logger = logging.getLogger(__name__)


DEFAULT_DYNAMIC_WRITE_STRING_ID = -1
DYNAMIC_TYPE_ID = -1
USE_CLASSNAME = 0
USE_CLASS_ID = 1
# preserve 0 as flag for class id not set in ClassInfo`
NO_CLASS_ID = 0
PYINT_CLASS_ID = 1
PYFLOAT_CLASS_ID = 2
PYBOOL_CLASS_ID = 3
STRING_CLASS_ID = 4
PICKLE_CLASS_ID = 5
PICKLE_STRONG_CACHE_CLASS_ID = 6
PICKLE_CACHE_CLASS_ID = 7
# `NOT_NULL_VALUE_FLAG` + `CLASS_ID << 1` in little-endian order
NOT_NULL_PYINT_FLAG = NOT_NULL_VALUE_FLAG & 0b11111111 | (PYINT_CLASS_ID << 9)
NOT_NULL_PYFLOAT_FLAG = NOT_NULL_VALUE_FLAG & 0b11111111 | (PYFLOAT_CLASS_ID << 9)
NOT_NULL_PYBOOL_FLAG = NOT_NULL_VALUE_FLAG & 0b11111111 | (PYBOOL_CLASS_ID << 9)
NOT_NULL_STRING_FLAG = NOT_NULL_VALUE_FLAG & 0b11111111 | (STRING_CLASS_ID << 9)
SMALL_STRING_THRESHOLD = 16


class ClassInfo:
__slots__ = (
"cls",
Expand All @@ -105,13 +98,13 @@ class ClassInfo:
)

def __init__(
self,
cls: type = None,
class_id: int = NO_CLASS_ID,
serializer: Serializer = None,
namespace_bytes=None,
typename_bytes=None,
dynamic_type: bool = False,
self,
cls: type = None,
class_id: int = NO_CLASS_ID,
serializer: Serializer = None,
namespace_bytes=None,
typename_bytes=None,
dynamic_type: bool = False,
):
self.cls = cls
self.class_id = class_id
Expand Down Expand Up @@ -143,9 +136,9 @@ class ClassResolver:
"_dynamic_written_metastr",
"_ns_type_to_classinfo",
"namespace_encoder",
"_namespace_decoder",
"_typename_encoder",
"_typename_decoder",
"namespace_decoder",
"typename_encoder",
"typename_decoder",
"require_registration",
"metastring_resolver",
"language",
Expand Down Expand Up @@ -186,12 +179,6 @@ def _initialize_py(self):
register(float, type_id=PYFLOAT_CLASS_ID, serializer=DoubleSerializer)
register(bool, type_id=PYBOOL_CLASS_ID, serializer=BooleanSerializer)
register(str, type_id=STRING_CLASS_ID, serializer=StringSerializer)
from pyfury import (
PickleCacheSerializer,
PickleStrongCacheSerializer,
PickleSerializer,
)

register(_PickleStub, type_id=PICKLE_CLASS_ID, serializer=PickleSerializer)
register(
PickleStrongCacheStub,
Expand Down Expand Up @@ -278,10 +265,10 @@ def _initialize_xlang(self):
# if pyarray are needed, one must annotate that value with XXXArrayType
# as a field of a struct.
for dtype, (
itemsize,
format,
ftype,
typeid,
itemsize,
format,
ftype,
typeid,
) in Numpy1DArraySerializer.dtypes_dict.items():
register(
ftype,
Expand All @@ -294,13 +281,13 @@ def _initialize_xlang(self):
register(dict, type_id=TypeId.MAP, serializer=MapSerializer)

def register_type(
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer=None,
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer=None,
):
self._register_type(
cls,
Expand All @@ -311,14 +298,14 @@ def register_type(
)

def _register_type(
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer=None,
internal=False,
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer=None,
internal=False,
):
"""Register class with given type id or typename. If typename is not None, it will be used for
cross-language serialization."""
Expand Down Expand Up @@ -351,14 +338,14 @@ def _register_type(
)

def _register_xtype(
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer=None,
internal=False,
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer=None,
internal=False,
):
if serializer is None:
if issubclass(cls, enum.Enum):
Expand All @@ -385,14 +372,14 @@ def _register_xtype(
)

def _register_pytype(
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer: Serializer = None,
internal: bool = False,
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer: Serializer = None,
internal: bool = False,
):
self.__register_type(
cls,
Expand All @@ -404,14 +391,14 @@ def _register_pytype(
)

def __register_type(
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer: Serializer = None,
internal: bool = False,
self,
cls: Union[type, TypeVar],
*,
type_id: int = None,
namespace: str = None,
typename: str = None,
serializer: Serializer = None,
internal: bool = False,
):
if not internal and serializer is None:
serializer = self._create_serializer(cls)
Expand Down Expand Up @@ -489,9 +476,9 @@ def _create_serializer(self, cls):
for clz in cls.__mro__:
class_info = self._classes_info.get(clz)
if (
class_info
and class_info.serializer
and class_info.serializer.support_subclass()
class_info
and class_info.serializer
and class_info.serializer.support_subclass()
):
serializer = type(class_info.serializer)(self.fury, cls)
break
Expand Down
Loading

0 comments on commit 473c796

Please sign in to comment.