1+ # Copyright 2023 JanusGraph-Python Authors
2+ #
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+ #
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+ #
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
14+
15+ from gremlin_python .structure .io .graphbinaryV1 import (
16+ _GraphBinaryTypeIO , StringIO , GraphBinaryReader , GraphBinaryWriter , DataType ,
17+ _make_packer ,
18+ uint64_pack , uint64_unpack , uint8_pack , uint8_unpack
19+ )
20+ from janusgraph_python .process .traversal import _JanusGraphP , RelationIdentifier
21+
22+ uint16_pack , uint16_unpack = _make_packer ('>H' )
23+ uint32_pack , uint32_unpack = _make_packer ('>I' )
24+
25+ class JanusGraphBinaryReader (GraphBinaryReader ):
26+ def __init__ (self ):
27+ # register JanusGraph-specific RelationIdentifier deserializer for custom type
28+ deserializer_map = {
29+ DataType .custom : JanusGraphCustomIO
30+ }
31+
32+ GraphBinaryReader .__init__ (self , deserializer_map )
33+
34+ class JanusGraphBinaryWriter (GraphBinaryWriter ):
35+ def __init__ (self ):
36+ # register JanusGraph-specific RelationIdentifier and text-predicate serializer
37+ serializer_map = [
38+ (RelationIdentifier , JanusGraphRelationIdentifierIO ),
39+ (_JanusGraphP , JanusGraphPSerializer )
40+ ]
41+ GraphBinaryWriter .__init__ (self , serializer_map )
42+
43+ class _JanusGraphBinaryTypeIO (_GraphBinaryTypeIO ):
44+ @classmethod
45+ def custom_type (cls , writer , to_extend , as_value = False ):
46+ if to_extend is None :
47+ to_extend = bytearray ()
48+
49+ # serializing the custom JanusGraph type
50+ # use the custom type code
51+ if not as_value : #- identifier does not like this if
52+ to_extend += uint8_pack (DataType .custom .value )
53+
54+ # add the name of the custom JanusGraph type
55+ StringIO .dictify (cls .graphbinary_type_name , writer , to_extend , True , False )
56+
57+ # add the id of the custom JanusGraph type
58+ to_extend += uint32_pack (cls .graphbinary_type_id )
59+
60+ # use the custom type code
61+ if not as_value :
62+ to_extend += uint8_pack (DataType .custom .value )
63+
64+ class JanusGraphPSerializer (_JanusGraphBinaryTypeIO ):
65+ graphbinary_type_id = 0x1002
66+ graphbinary_type_name = "janusgraph.P"
67+ python_type = _JanusGraphP
68+
69+ @classmethod
70+ def dictify (cls , obj , writer , to_extend , as_value = False , nullable = True ):
71+ cls .custom_type (writer , to_extend , as_value )
72+
73+ # serializing the custom JanusGraph operator
74+ StringIO .dictify (obj .operator , writer , to_extend , True , False )
75+
76+ # serialize the value
77+ writer .to_dict (obj .value , to_extend )
78+
79+ return to_extend
80+
81+ class JanusGraphRelationIdentifierIO (_JanusGraphBinaryTypeIO ):
82+ graphbinary_type_id = 0x1001
83+ graphbinary_type_name = "janusgraph.RelationIdentifier"
84+ python_type = RelationIdentifier
85+
86+ long_marker = 0
87+ string_marker = 1
88+
89+ @classmethod
90+ def _write_string (cls , string , writer , to_extend ):
91+ b = bytearray ()
92+ b .extend (map (ord , string ))
93+ b [- 1 ] |= 0x80 # add end marker to the last character
94+ to_extend += b
95+
96+ @classmethod
97+ def dictify (cls , obj , writer , to_extend , as_value = False , nullable = True ):
98+ cls .custom_type (writer , to_extend , as_value )
99+
100+ if isinstance (obj .out_vertex_id , int ):
101+ to_extend += uint8_pack (cls .long_marker )
102+ to_extend += uint64_pack (obj .out_vertex_id )
103+ else :
104+ to_extend += uint8_pack (cls .string_marker )
105+ cls ._write_string (obj .out_vertex_id , writer , to_extend )
106+
107+ to_extend += uint64_pack (obj .type_id )
108+ to_extend += uint64_pack (obj .relation_id )
109+
110+ if obj .in_vertex_id is None :
111+ to_extend += uint8_pack (cls .long_marker )
112+ to_extend += uint64_pack (0 )
113+ elif isinstance (obj .in_vertex_id , int ):
114+ to_extend += uint8_pack (cls .long_marker )
115+ to_extend += uint64_pack (obj .in_vertex_id )
116+ else :
117+ to_extend += uint8_pack (cls .string_marker )
118+ cls ._write_string (obj .in_vertex_id , writer , to_extend )
119+
120+ return to_extend
121+
122+ @classmethod
123+ def _read_string (cls , buff ):
124+ final_string = ""
125+ while True :
126+ c = 0xFF & uint8_unpack (buff .read (1 ))
127+ final_string += chr (c & 0x7F )
128+ if c & 0x80 > 0 :
129+ break
130+
131+ return final_string
132+
133+ @classmethod
134+ def objectify (cls , b , r ):
135+ if uint8_unpack (b .read (1 )) != DataType .custom .value :
136+ raise Exception ("Unexpected type while deserializing JanusGraph RelationIdentifier" )
137+
138+ # read the next byte that shows if the out vertex id is string or long
139+ out_vertex_id_marker = uint8_unpack (b .read (1 ))
140+
141+ if out_vertex_id_marker == cls .string_marker :
142+ out_vertex_id = cls ._read_string (b )
143+ else :
144+ out_vertex_id = uint64_unpack (b .read (8 ))
145+
146+ type_id = uint64_unpack (b .read (8 ))
147+ relation_id = uint64_unpack (b .read (8 ))
148+
149+ in_vertex_id_marker = uint8_unpack (b .read (1 ))
150+ if in_vertex_id_marker == cls .string_marker :
151+ in_vertex_id = cls ._read_string (b )
152+ else :
153+ in_vertex_id = uint64_unpack (b .read (8 ))
154+ if in_vertex_id == 0 :
155+ in_vertex_id = None
156+
157+ return RelationIdentifier .from_ids (out_vertex_id , type_id , relation_id , in_vertex_id )
158+
159+ class JanusGraphCustomIO (_GraphBinaryTypeIO ):
160+ # list of JanusGraph custom types with their type_id, type_name and class for deserialization
161+ io_registry = {
162+ 0x1001 : ("janusgraph.RelationIdentifier" , JanusGraphRelationIdentifierIO )
163+ }
164+
165+ @classmethod
166+ def objectify (cls , buff , reader , nullable = True ):
167+ return cls .is_null (buff , reader , cls ._read_data , nullable )
168+
169+ @classmethod
170+ def _read_data (cls , b , r ):
171+ # check if first byte is custom byte notation
172+ if uint8_unpack (b .read (1 )) != DataType .custom .value :
173+ return None
174+
175+ # get the custom type name length
176+ custom_type_name_length = uint16_unpack (b .read (2 ))
177+ custom_type_name = b .read (custom_type_name_length ).decode ()
178+
179+ # read the custom type id
180+ custom_type_id = uint32_unpack (b .read (4 ))
181+
182+ custom_serializer = cls .io_registry .get (custom_type_id )
183+ if not custom_serializer :
184+ raise NotImplementedError (f"No JanusGraph serializer found for type with id: { custom_type_id } )
185+
186+ # check the type name
187+ if custom_serializer [0 ] != custom_type_name :
188+ raise NotImplementedError (f"No JanusGraph serializer found for type with name: { custom_type_name } )
189+
190+ return custom_serializer [1 ].objectify (b , r )
0 commit comments