tuple.py

# FoundationDB Python API
# Copyright (c) 2012 FoundationDB, LLC

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import struct
import sys
from bisect import bisect_left

if sys.version < '3':
    _integer_types = (int, long,)
    _chr = chr
else:
    _integer_types = (int,)
    unicode = str
    def _chr(x):
        return bytes((x,))

_size_limits = tuple( (1 << (i*8))-1 for i in range(9) )

def _find_terminator( v, pos ):
    # Finds the start of the next terminator [\x00]![\xff] or the end of v
    while True:
        pos = v.find(b'\x00', pos)
        if pos < 0:
            return len(v)
        if pos+1 == len(v) or v[pos+1:pos+2] != b'\xff':
            return pos
        pos += 2

def _decode(v, pos):
    if isinstance(v[pos], _integer_types):
        code = v[pos]
    else:
        code = ord(v[pos])
    if code == 0:
        return None, pos+1
    elif code == 1:
        end = _find_terminator(v, pos+1)
        return v[pos+1:end].replace(b"\x00\xFF", b"\x00"), end+1
    elif code == 2:
        end = _find_terminator(v, pos+1)
        return v[pos+1:end].replace(b"\x00\xFF", b"\x00").decode("utf-8"), end+1
    elif code >= 20 and code <= 28:
        n = code - 20
        end = pos + 1 + n
        return struct.unpack(">Q", b'\x00'*(8-n) + v[pos+1:end])[0], end
    elif code >= 12 and code < 20:
        n = 20 - code
        end = pos + 1 + n
        return struct.unpack(">Q", b'\x00'*(8-n) + v[pos+1:end])[0]-_size_limits[n], end
    else:
        print(code)
        raise ValueError("Unknown data type in DB: " + repr(v))

def _encode(value):
    # returns [code][data] (code != 0xFF)
    # encoded values are self-terminating
    # sorting need to work too!
    if value == None:  # ==, not is, because some fdb.impl.Value are equal to None
        return b'\x00'
    elif isinstance(value, bytes): # also gets non-None fdb.impl.Value
        return b'\x01' + value.replace(b'\x00', b'\x00\xFF') + b'\x00'
    elif isinstance(value, unicode):
        return b'\x02' + value.encode('utf-8').replace(b'\x00', b'\x00\xFF') + b'\x00'
    elif isinstance(value, _integer_types):
        if value == 0:
            return b'\x14'
        elif value > 0:
            n = bisect_left( _size_limits, value )
            return _chr(20 + n) + struct.pack( ">Q", value )[-n:]
        else:
            n = bisect_left( _size_limits, -value )
            maxv = _size_limits[n]
            return _chr(20 - n) + struct.pack( ">Q", maxv+value)[-n:]
    else:
        raise ValueError("Unsupported data type: " + str(type(value)))

# packs the specified tuple into a key
def pack(t):
    if not isinstance(t, tuple):
        raise Exception("fdbtuple pack() expects a tuple, got a " + str(type(t)))
    return b''.join([_encode(x) for x in t])

# unpacks the specified key into a tuple
def unpack(key):
    pos = 0
    res = []
    while pos < len(key):
        r, pos = _decode(key, pos)
        res.append(r)
    return tuple(res)

_range = range
def range(t):
    """Returns a slice of keys that includes all tuples of greater
    length than the specified tuple that that start with the
    specified elements.

    e.g. range(('a', 'b')) includes all tuples ('a', 'b', ...)"""

    if not isinstance(t, tuple):
        raise Exception("fdbtuple range() expects a tuple, got a " + str(type(t)))

    p = pack(t)
    return slice(
        p+b'\x00',
        p+b'\xff')