* Make StringStore.__getitem__ accept unicode-typed keys.

This commit is contained in:
Matthew Honnibal 2014-12-03 01:33:20 +11:00
parent 522bb0346e
commit 14097311ae
3 changed files with 17 additions and 9 deletions

View File

@ -13,7 +13,7 @@ cdef struct Utf8Str:
cdef class StringStore: cdef class StringStore:
cdef Pool mem cdef Pool mem
cdef PreshMap table cdef PreshMap _map
cdef Utf8Str* strings cdef Utf8Str* strings
cdef int size cdef int size
cdef int _resize_at cdef int _resize_at

View File

@ -8,7 +8,7 @@ SEPARATOR = '\n|-SEP-|\n'
cdef class StringStore: cdef class StringStore:
def __init__(self): def __init__(self):
self.mem = Pool() self.mem = Pool()
self.table = PreshMap() self._map = PreshMap()
self._resize_at = 10000 self._resize_at = 10000
self.strings = <Utf8Str*>self.mem.alloc(self._resize_at, sizeof(Utf8Str)) self.strings = <Utf8Str*>self.mem.alloc(self._resize_at, sizeof(Utf8Str))
self.size = 1 self.size = 1
@ -17,17 +17,21 @@ cdef class StringStore:
def __get__(self): def __get__(self):
return self.size-1 return self.size-1
def __getitem__(self, string_or_id): def __getitem__(self, object string_or_id):
cdef bytes byte_string cdef bytes byte_string
cdef Utf8Str* utf8str cdef Utf8Str* utf8str
if type(string_or_id) == int or type(string_or_id) == long: if isinstance(string_or_id, int):
if string_or_id < 1 or string_or_id >= self.size: if string_or_id < 1 or string_or_id >= self.size:
raise IndexError(string_or_id) raise IndexError(string_or_id)
utf8str = &self.strings[<int>string_or_id] utf8str = &self.strings[<int>string_or_id]
return utf8str.chars[:utf8str.length] return utf8str.chars[:utf8str.length]
elif type(string_or_id) == bytes: elif isinstance(string_or_id, bytes):
utf8str = self.intern(<char*>string_or_id, len(string_or_id)) utf8str = self.intern(<char*>string_or_id, len(string_or_id))
return utf8str.i return utf8str.i
elif isinstance(string_or_id, unicode):
byte_string = string_or_id.encode('utf8')
utf8str = self.intern(<char*>byte_string, len(byte_string))
return utf8str.i
else: else:
raise TypeError(type(string_or_id)) raise TypeError(type(string_or_id))
@ -36,7 +40,7 @@ cdef class StringStore:
# slot 0 to simplify the code, because it doesn't matter. # slot 0 to simplify the code, because it doesn't matter.
assert length != 0 assert length != 0
cdef hash_t key = hash64(chars, length * sizeof(char), 0) cdef hash_t key = hash64(chars, length * sizeof(char), 0)
cdef void* value = self.table.get(key) cdef void* value = self._map.get(key)
cdef size_t i cdef size_t i
if value == NULL: if value == NULL:
if self.size == self._resize_at: if self.size == self._resize_at:
@ -48,7 +52,7 @@ cdef class StringStore:
self.strings[i].chars = <char*>self.mem.alloc(length, sizeof(char)) self.strings[i].chars = <char*>self.mem.alloc(length, sizeof(char))
memcpy(self.strings[i].chars, chars, length) memcpy(self.strings[i].chars, chars, length)
self.strings[i].length = length self.strings[i].length = length
self.table.set(key, <void*>self.size) self._map.set(key, <void*>self.size)
self.size += 1 self.size += 1
else: else:
i = <size_t>value i = <size_t>value

View File

@ -19,8 +19,12 @@ def test_save_bytes(sstore):
def test_save_unicode(sstore): def test_save_unicode(sstore):
with pytest.raises(TypeError): Hello_i = sstore[u'Hello']
A_i = sstore['A'] assert Hello_i == 1
assert sstore[u'Hello'] == 1
assert sstore[u'goodbye'] != Hello_i
assert sstore[u'hello'] != Hello_i
assert Hello_i == 1
def test_zero_id(sstore): def test_zero_id(sstore):