1import struct 2from math import floor 3from ...python.functools import Composition as compose 4from ...python.itertools import interlace 5from ...python.structlib import \ 6 short_pack, short_unpack, \ 7 ulong_pack, ulong_unpack, \ 8 long_pack, long_unpack, \ 9 double_pack, double_unpack, \ 10 longlong_pack, longlong_unpack, \ 11 float_pack, float_unpack, \ 12 LH_pack, LH_unpack, \ 13 dl_pack, dl_unpack, \ 14 dll_pack, dll_unpack, \ 15 ql_pack, ql_unpack, \ 16 qll_pack, qll_unpack, \ 17 llL_pack, llL_unpack, \ 18 dd_pack, dd_unpack, \ 19 ddd_pack, ddd_unpack, \ 20 dddd_pack, dddd_unpack, \ 21 hhhh_pack, hhhh_unpack 22 23oid_pack = cid_pack = xid_pack = ulong_pack 24oid_unpack = cid_unpack = xid_unpack = ulong_unpack 25tid_pack, tid_unpack = LH_pack, LH_unpack 26 27# geometry types 28point_pack, point_unpack = dd_pack, dd_unpack 29circle_pack, circle_unpack = ddd_pack, ddd_unpack 30lseg_pack = box_pack = dddd_pack 31lseg_unpack = box_unpack = dddd_unpack 32 33null_sequence = b'\xff\xff\xff\xff' 34string_format = b'\x00\x00' 35binary_format = b'\x00\x01' 36 37def numeric_pack(data, hhhh_pack = hhhh_pack, pack = struct.pack, len = len): 38 return hhhh_pack(data[0]) + pack("!%dh"%(len(data[1]),), *data[1]) 39 40def numeric_unpack(data, hhhh_unpack = hhhh_unpack, unpack = struct.unpack, len = len): 41 return (hhhh_unpack(data[:8]), unpack("!8x%dh"%((len(data)-8) // 2,), data)) 42 43def path_pack(data, pack = struct.pack, len = len): 44 """ 45 Given a sequence of point data, pack it into a path's serialized form. 46 47 [px1, py1, px2, py2, ...] 48 49 Must be an even number of numbers. 50 """ 51 return pack("!l%dd" %(len(data),), len(data), *data) 52 53def path_unpack(data, long_unpack = long_unpack, unpack = struct.unpack): 54 """ 55 Unpack a path's serialized form into a sequence of point data: 56 57 [px1, py1, px2, py2, ...] 58 59 Should be an even number of numbers. 60 """ 61 return unpack("!4x%dd" %(long_unpack(data[:4]),), data) 62polygon_pack, polygon_unpack = path_pack, path_unpack 63 64## 65# Binary representations of infinity for datetimes. 66time_infinity = b'\x7f\xf0\x00\x00\x00\x00\x00\x00' 67time_negative_infinity = b'\xff\xf0\x00\x00\x00\x00\x00\x00' 68time64_infinity = b'\x7f\xff\xff\xff\xff\xff\xff\xff' 69time64_negative_infinity = b'\x80\x00\x00\x00\x00\x00\x00\x00' 70date_infinity = b'\x7f\xff\xff\xff' 71date_negative_infinity = b'\x80\x00\x00\x00' 72 73# time types 74date_pack, date_unpack = long_pack, long_unpack 75 76def mktimetuple(ts, floor = floor): 77 'make a pair of (seconds, microseconds) out of the given double' 78 seconds = floor(ts) 79 return (int(seconds), int(1000000 * (ts - seconds))) 80 81def mktimetuple64(ts, divmod = divmod): 82 'make a pair of (seconds, microseconds) out of the given long' 83 return divmod(ts, 1000000) 84 85def mktime(seconds_ms, float = float): 86 'make a double out of the pair of (seconds, microseconds)' 87 return float(seconds_ms[0]) + (seconds_ms[1] / 1000000.0) 88 89def mktime64(seconds_ms): 90 'make an integer out of the pair of (seconds, microseconds)' 91 return seconds_ms[0] * 1000000 + seconds_ms[1] 92 93# takes a pair, (seconds, microseconds) 94time_pack = compose((mktime, double_pack)) 95time_unpack = compose((double_unpack, mktimetuple)) 96 97def interval_pack(m_d_timetup, mktime = mktime, dll_pack = dll_pack): 98 """ 99 Given a triple, (month, day, (seconds, microseconds)), serialize it for 100 transport. 101 """ 102 (month, day, timetup) = m_d_timetup 103 return dll_pack((mktime(timetup), day, month)) 104 105def interval_unpack(data, dll_unpack = dll_unpack, mktimetuple = mktimetuple): 106 """ 107 Given a serialized interval, '{month}{day}{time}', yield the triple: 108 109 (month, day, (seconds, microseconds)) 110 """ 111 tim, day, month = dll_unpack(data) 112 return (month, day, mktimetuple(tim)) 113 114def interval_noday_pack(month_day_timetup, dl_pack = dl_pack, mktime = mktime): 115 """ 116 Given a triple, (month, day, (seconds, microseconds)), return the serialized 117 form that does not have an individual day component. 118 119 There is no day component, so if day is non-zero, it will be converted to 120 seconds and subsequently added to the seconds. 121 """ 122 (month, day, timetup) = month_day_timetup 123 if day: 124 timetup = (timetup[0] + (day * 24 * 60 * 60), timetup[1]) 125 return dl_pack((mktime(timetup), month)) 126 127def interval_noday_unpack(data, dl_unpack = dl_unpack, mktimetuple = mktimetuple): 128 """ 129 Given a serialized interval without a day component, return the triple: 130 131 (month, day(always zero), (seconds, microseconds)) 132 """ 133 tim, month = dl_unpack(data) 134 return (month, 0, mktimetuple(tim)) 135 136def time64_pack(data, mktime64 = mktime64, longlong_pack = longlong_pack): 137 return longlong_pack(mktime64(data)) 138def time64_unpack(data, longlong_unpack = longlong_unpack, mktimetuple64 = mktimetuple64): 139 return mktimetuple64(longlong_unpack(data)) 140 141def interval64_pack(m_d_timetup, qll_pack = qll_pack, mktime64 = mktime64): 142 """ 143 Given a triple, (month, day, (seconds, microseconds)), return the serialized 144 data using a quad-word for the (seconds, microseconds) tuple. 145 """ 146 (month, day, timetup) = m_d_timetup 147 return qll_pack((mktime64(timetup), day, month)) 148 149def interval64_unpack(data, qll_unpack = qll_unpack, mktimetuple = mktimetuple): 150 """ 151 Unpack an interval containing a quad-word into a triple: 152 153 (month, day, (seconds, microseconds)) 154 """ 155 tim, day, month = qll_unpack(data) 156 return (month, day, mktimetuple64(tim)) 157 158def interval64_noday_pack(m_d_timetup, ql_pack = ql_pack, mktime64 = mktime64): 159 """ 160 Pack an interval without a day component and using a quad-word for second 161 representation. 162 163 There is no day component, so if day is non-zero, it will be converted to 164 seconds and subsequently added to the seconds. 165 """ 166 (month, day, timetup) = m_d_timetup 167 if day: 168 timetup = (timetup[0] + (day * 24 * 60 * 60), timetup[1]) 169 return ql_pack((mktime64(timetup), month)) 170 171def interval64_noday_unpack(data, ql_unpack = ql_unpack, mktimetuple64 = mktimetuple64): 172 """ 173 Unpack a ``noday`` quad-word based interval. Returns a triple: 174 175 (month, day(always zero), (seconds, microseconds)) 176 """ 177 tim, month = ql_unpack(data) 178 return (month, 0, mktimetuple64(tim)) 179 180def timetz_pack(timetup_tz, dl_pack = dl_pack, mktime = mktime): 181 """ 182 Pack a time; offset from beginning of the day and timezone offset. 183 184 Given a pair, ((seconds, microseconds), timezone_offset), pack it into its 185 serialized form: "!dl". 186 """ 187 (timetup, tz_offset) = timetup_tz 188 return dl_pack((mktime(timetup), tz_offset)) 189 190def timetz_unpack(data, dl_unpack = dl_unpack, mktimetuple = mktimetuple): 191 """ 192 Given serialized time data, unpack it into a pair: 193 194 ((seconds, microseconds), timezone_offset). 195 """ 196 ts, tz = dl_unpack(data) 197 return (mktimetuple(ts), tz) 198 199def timetz64_pack(timetup_tz, ql_pack = ql_pack, mktime64 = mktime64): 200 """ 201 Pack a time; offset from beginning of the day and timezone offset. 202 203 Given a pair, ((seconds, microseconds), timezone_offset), pack it into its 204 serialized form using a long long: "!ql". 205 """ 206 (timetup, tz_offset) = timetup_tz 207 return ql_pack((mktime64(timetup), tz_offset)) 208 209def timetz64_unpack(data, ql_unpack = ql_unpack, mktimetuple64 = mktimetuple64): 210 """ 211 Given "long long" serialized time data, "ql", unpack it into a pair: 212 213 ((seconds, microseconds), timezone_offset) 214 """ 215 ts, tz = ql_unpack(data) 216 return (mktimetuple64(ts), tz) 217 218# oidvectors are 128 bytes, so pack the number of Oids in self 219# and justify that to 128 by padding with \x00. 220def oidvector_pack(seq, pack = struct.pack): 221 """ 222 Given a sequence of Oids, pack them into the serialized form. 223 224 An oidvector is a type used by the PostgreSQL catalog. 225 """ 226 return pack("!%dL"%(len(seq),), *seq).ljust(128, '\x00') 227 228def oidvector_unpack(data, unpack = struct.unpack): 229 """ 230 Given a serialized oidvector(32 longs), unpack it into a list of unsigned integers. 231 232 An int2vector is a type used by the PostgreSQL catalog. 233 """ 234 return unpack("!32L", data) 235 236def int2vector_pack(seq, pack = struct.pack): 237 """ 238 Given a sequence of integers, pack them into the serialized form. 239 240 An int2vector is a type used by the PostgreSQL catalog. 241 """ 242 return pack("!%dh"%(len(seq),), *seq).ljust(64, '\x00') 243 244def int2vector_unpack(data, unpack = struct.unpack): 245 """ 246 Given a serialized int2vector, unpack it into a list of integers. 247 248 An int2vector is a type used by the PostgreSQL catalog. 249 """ 250 return unpack("!32h", data) 251 252def varbit_pack(bits_data, long_pack = long_pack): 253 r""" 254 Given a pair, serialize the varbit. 255 256 # (number of bits, data) 257 >>> varbit_pack((1, '\x00')) 258 b'\x00\x00\x00\x01\x00' 259 """ 260 return long_pack(bits_data[0]) + bits_data[1] 261 262def varbit_unpack(data, long_unpack = long_unpack): 263 """ 264 Given ``varbit`` data, unpack it into a pair: 265 266 (bits, data) 267 268 Where bits are the total number of bits in data (bytes). 269 """ 270 return long_unpack(data[0:4]), data[4:] 271 272def net_pack(triple, 273 # Map PGSQL src/include/utils/inet.h to IP version number. 274 fmap = { 275 4: 2, 276 6: 3, 277 }, 278 len = len, 279): 280 """ 281 net_pack((family, mask, data)) 282 283 Pack Postgres' inet/cidr data structure. 284 """ 285 family, mask, data = triple 286 return bytes((fmap[family], mask or 0, 0 if mask is None else 1, len(data))) + data 287 288def net_unpack(data, 289 # Map IP version number to PGSQL src/include/utils/inet.h. 290 fmap = { 291 2: 4, 292 3: 6, 293 } 294): 295 """ 296 net_unpack(data) 297 298 Unpack Postgres' inet/cidr data structure. 299 """ 300 family, mask, is_cidr, size = data[:4] 301 return (fmap[family], mask, data[4:]) 302 303def macaddr_pack(data, bytes = bytes): 304 """ 305 Pack a MAC address 306 307 Format found in PGSQL src/backend/utils/adt/mac.c, and PGSQL Manual types 308 """ 309 # Accept all possible PGSQL Macaddr formats as in manual 310 # Oh for sscanf() as we could just copy PGSQL C in src/util/adt/mac.c 311 colon_parts = data.split(':') 312 dash_parts = data.split('-') 313 dot_parts = data.split('.') 314 if len(colon_parts) == 6: 315 mac_parts = colon_parts 316 elif len(dash_parts) == 6: 317 mac_parts = dash_parts 318 elif len(colon_parts) == 2: 319 mac_parts = [colon_parts[0][:2], colon_parts[0][2:4], colon_parts[0][4:], 320 colon_parts[1][:2], colon_parts[1][2:4], colon_parts[1][4:]] 321 elif len(dash_parts) == 2: 322 mac_parts = [dash_parts[0][:2], dash_parts[0][2:4], dash_parts[0][4:], 323 dash_parts[1][:2], dash_parts[1][2:4], dash_parts[1][4:]] 324 elif len(dot_parts) == 3: 325 mac_parts = [dot_parts[0][:2], dot_parts[0][2:], dot_parts[1][:2], 326 dot_parts[1][2:], dot_parts[2][:2], dot_parts[2][2:]] 327 elif len(colon_parts) == 1: 328 mac_parts = [data[:2], data[2:4], data[4:6], data[6:8], data[8:10], data[10:]] 329 else: 330 raise ValueError('data string cannot be parsed to bytes') 331 if len(mac_parts) != 6 and len(mac_parts[-1]) != 2: 332 raise ValueError('data string cannot be parsed to bytes') 333 return bytes([int(p, 16) for p in mac_parts]) 334 335def macaddr_unpack(data): 336 """ 337 Unpack a MAC address 338 339 Format found in PGSQL src/backend/utils/adt/mac.c 340 """ 341 # This is easy, just go for standard macaddr format, 342 # just like PGSQL in src/util/adt/mac.c macaddr_out() 343 if len(data) != 6: 344 raise ValueError('macaddr has incorrect length') 345 return ("%02x:%02x:%02x:%02x:%02x:%02x" % tuple(data)) 346 347def record_unpack(data, 348 long_unpack = long_unpack, 349 oid_unpack = oid_unpack, 350 null_sequence = null_sequence, 351 len = len): 352 """ 353 Given serialized record data, return a tuple of tuples of type Oids and 354 attributes. 355 """ 356 columns = long_unpack(data) 357 offset = 4 358 359 for x in range(columns): 360 typid = oid_unpack(data[offset:offset+4]) 361 offset += 4 362 363 if data[offset:offset+4] == null_sequence: 364 att = None 365 offset += 4 366 else: 367 size = long_unpack(data[offset:offset+4]) 368 offset += 4 369 att = data[offset:offset + size] 370 if size < -1 or len(att) != size: 371 raise ValueError("insufficient data left in message") 372 offset += size 373 yield (typid, att) 374 375 if len(data) - offset != 0: 376 raise ValueError("extra data, %d octets, at end of record" %(len(data),)) 377 378def record_pack(seq, 379 long_pack = long_pack, 380 oid_pack = oid_pack, 381 null_sequence = null_sequence): 382 """ 383 pack a record given an iterable of (type_oid, data) pairs. 384 """ 385 return long_pack(len(seq)) + b''.join([ 386 # typid + (null_seq or data) 387 oid_pack(x) + (y is None and null_sequence or (long_pack(len(y)) + y)) 388 for x, y in seq 389 ]) 390 391def elements_pack(elements, 392 null_sequence = null_sequence, 393 long_pack = long_pack, len = len 394): 395 """ 396 Pack the elements for containment within a serialized array. 397 398 This is used by array_pack. 399 """ 400 for x in elements: 401 if x is None: 402 yield null_sequence 403 else: 404 yield long_pack(len(x)) 405 yield x 406 407def array_pack(array_data, 408 llL_pack = llL_pack, 409 len = len, 410 long_pack = long_pack, 411 interlace = interlace 412): 413 """ 414 Pack a raw array. A raw array consists of flags, type oid, sequence of lower 415 and upper bounds, and an iterable of already serialized element data: 416 417 (0, element type oid, (lower bounds, upper bounds, ...), iterable of element_data) 418 419 The lower bounds and upper bounds specifies boundaries of the dimension. So the length 420 of the boundaries sequence is two times the number of dimensions that the array has. 421 422 array_pack((flags, type_id, dims, lowers, element_data)) 423 424 The format of ``lower_upper_bounds`` is a sequence of lower bounds and upper 425 bounds. First lower then upper inlined within the sequence: 426 427 [lower, upper, lower, upper] 428 429 The above array `dlb` has two dimensions. The lower and upper bounds of the 430 first dimension is defined by the first two elements in the sequence. The 431 second dimension is then defined by the last two elements in the sequence. 432 """ 433 (flags, typid, dims, lbs, elements) = array_data 434 return llL_pack((len(dims), flags, typid)) + \ 435 b''.join(map(long_pack, interlace(dims, lbs))) + \ 436 b''.join(elements_pack(elements)) 437 438def elements_unpack(data, offset, 439 long_unpack = long_unpack, 440 null_sequence = null_sequence): 441 """ 442 Unpack the serialized elements of an array into a list. 443 444 This is used by array_unpack. 445 """ 446 data_len = len(data) 447 while offset < data_len: 448 lend = data[offset:offset+4] 449 offset += 4 450 if lend == null_sequence: 451 yield None 452 else: 453 sizeof_el = long_unpack(lend) 454 yield data[offset:offset+sizeof_el] 455 offset += sizeof_el 456 457def array_unpack(data, 458 llL_unpack = llL_unpack, 459 unpack = struct.unpack_from, 460 long_unpack = long_unpack 461): 462 """ 463 Given a serialized array, unpack it into a tuple: 464 465 (flags, typid, (dims, lower bounds, ...), [elements]) 466 """ 467 ndim, flags, typid = llL_unpack(data) 468 if ndim < 0: 469 raise ValueError("invalid number of dimensions: %d" %(ndim,)) 470 # "ndim" number of pairs of longs 471 end = (4 * 2 * ndim) + 12 472 # Dimensions and lower bounds; split the two early. 473 #dlb = unpack("!%dl"%(2 * ndim,), data, 12) 474 dims = [long_unpack(data[x:x+4]) for x in range(12, end, 8)] 475 lbs = [long_unpack(data[x:x+4]) for x in range(16, end, 8)] 476 return (flags, typid, dims, lbs, elements_unpack(data, end)) 477