polyfile-weave 0.5.5__py3-none-any.whl

polyfile/search.py

@@ -0,0 +1,310 @@
+from collections import deque
+import collections.abc
+from typing import IO, Mapping, Sequence, Union
+
+from . import serialization
+
+
+@serialization.serializable
+class TrieNode:
+    def __init__(self, value=None, sources=None, _children=None):
+        if _children is None:
+            self._children: Mapping[object, TrieNode] = {}
+        else:
+            self._children = _children
+        self.value = value
+        if sources is not None:
+            self._sources = set(sources)
+        else:
+            self._sources = set()
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(value={self.value!r}, sources={self.sources!r}, _children={self._children!r})"
+
+    def __len__(self):
+        return len(self._children)
+
+    def __iter__(self):
+        return iter(self._children.values())
+
+    def __hash__(self):
+        return hash(self.value)
+
+    def __eq__(self, other):
+        return self.value == other.value
+
+    def __ne__(self, other):
+        return not (self == other)
+
+    def __getitem__(self, key):
+        return self._children[key]
+
+    def __contains__(self, value):
+        if not isinstance(value, collections.abc.Sequence):
+            first, remainder, n = value, (), 1
+        else:
+            first, remainder, n = self._car_cdr_len(value)
+        if n == 1:
+            return first in self._children
+        else:
+            return self._find(first, remainder, n)
+
+    @staticmethod
+    def _car_cdr_len(sequence):
+        n = len(sequence)
+        if n == 0:
+            first = None
+        else:
+            first = sequence[0]
+        return first, sequence[1:], n
+
+    def _find(self, first, remainder, n):
+        if n == 0:
+            return len(self._sources) > 0
+        if first not in self._children:
+            return False
+        return self[first]._find(*self._car_cdr_len(remainder))
+
+    def find(self, key):
+        if isinstance(key, collections.abc.Sequence):
+            return self._find(*self._car_cdr_len(key))
+        else:
+            return self._find(key, (), 1)
+
+    @property
+    def children(self):
+        return dict(self._children)
+
+    @property
+    def sources(self):
+        return frozenset(self._sources)
+
+    def _add_child(self, value, sources=None):
+        new_child = TrieNode(value, sources)
+        self._children[value] = new_child
+        return new_child
+
+    def _add(self, sequence, source):
+        node = self
+        while True:
+            first, sequence, n = self._car_cdr_len(sequence)
+            if n == 0:
+                break
+            if first in node:
+                node = node[first]
+            else:
+                node = node._add_child(first)
+        node._sources.add(source)
+        return node
+
+    def add(self, sequence, source=None):
+        if source is None:
+            source = sequence
+        return self._add(sequence, source)
+
+    def find_prefix(self, prefix):
+        first, remainder, n = self._car_cdr_len(prefix)
+        if n == 0:
+            yield from iter(self._sources)
+            for child in self:
+                yield from child.find_prefix(prefix)
+        else:
+            if first in self._children:
+                yield from self[first].find_prefix(remainder)
+
+    def bfs(self):
+        queue = deque([self])
+        while queue:
+            head = queue.popleft()
+            yield head
+            queue.extend(head._children.values())
+
+    def dfs(self):
+        stack = [self]
+        visited = set()
+        while stack:
+            tail = stack.pop()
+            yield tail
+            children = tail._children.values()
+            stack.extend(child for child in children if id(child) not in visited)
+            visited |= set(id(c) for c in children)
+
+    def serialize(self):
+        return self.value, self._sources, self._children
+
+
+@serialization.serializable
+class ACNode(TrieNode):
+    """A data structure for implementing the Aho-Corasick multi-string matching algorithm"""
+    def __init__(self, value=None, sources=None, _children=None, parent=None, _fall=None):
+        super().__init__(value=value, sources=sources, _children=_children)
+        self.parent = parent
+        self.fall = _fall
+
+    def serialize(self):
+        return super().serialize() + (self.parent, self.fall)
+
+    def __repr__(self):
+        if self.fall is not self:
+            return f"{self.__class__.__name__}(value={self.value!r}, sources={self.sources!r}, _children={self._children!r}), parent={self.parent!r}, _fall={self.fall!r}"
+        else:
+            return f"{self.__class__.__name__}(value={self.value!r}, sources={self.sources!r}, _children={self._children!r}), parent={self.parent!r}, _fall=self"
+
+    def _add_child(self, value, sources=None):
+        new_child = ACNode(value, sources, parent=self)
+        self._children[value] = new_child
+        return new_child
+
+    def finalize(self):
+        self.fall = self
+        for n in self.bfs():
+            if n is self:
+                continue
+            new_fall = n.parent.fall
+            while n.value not in new_fall and new_fall is not self:
+                new_fall = new_fall.fall
+            if n.value not in new_fall:
+                # there is no suffix
+                n.fall = self
+            else:
+                n.fall = new_fall[n.value]
+                if n.fall is n:
+                    n.fall = self
+
+    def to_dot(self, include_falls=False):
+        """Returns a Graphviz/Dot representation of this Trie"""
+        dot = "digraph G {\n"
+        node_ids = {}
+        falls = {}
+        for node in self.dfs():
+            assert node not in node_ids
+            nid = len(node_ids)
+            node_ids[id(node)] = nid
+            dot += f"    node{nid}"
+            if node.value is None:
+                dot += f"[label=\"Root\"]"
+            else:
+                if node.value == ord('"'):
+                    c = '\\"'
+                elif node.value == ord('\\'):
+                    c = '\\\\'
+                elif 32 <= node.value <= 126:
+                    c = chr(node.value)
+                else:
+                    c = f"\\\\x{hex(node.value)[2:]}"
+                dot += f"[label=\"{c}\"]"
+            dot += ";\n"
+            if node.parent is not None:
+                dot += f"    node{node_ids[id(node.parent)]} -> node{nid};\n"
+            if include_falls and node.fall is not None and node.fall is not node:
+                falls[id(node)] = id(node.fall)
+
+        for nodeid, fallid in falls.items():
+            dot += f"    node{node_ids[nodeid]} -> node{node_ids[fallid]} [style=dashed,label=\"fall\"];\n"
+        dot += "}\n"
+        return dot
+
+
+class MultiSequenceSearch:
+    """A datastructure for efficiently searching a sequence for multiple strings"""
+    def __init__(self, *sequences_to_find):
+        self.trie = ACNode()
+        for seq in sequences_to_find:
+            self.trie.add(seq)
+        self.trie.finalize()
+
+    def save(self, output_stream: IO):
+        serialization.dump(self.trie, output_stream)
+
+    @staticmethod
+    def load(input_stream: IO):
+        mss = MultiSequenceSearch()
+        mss.trie = serialization.load(input_stream)
+        print(mss.trie)
+        exit(0)
+        return mss
+
+    def search(self, source_sequence: Union[Sequence, IO]):
+        """The Aho-Corasick Algorithm"""
+        if hasattr(source_sequence, 'read'):
+            def iterator():
+                while True:
+                    b = source_sequence.read(1)
+                    if not b:
+                        return
+                    yield b[0]
+        else:
+            def iterator():
+                return iter(source_sequence)
+
+        state = self.trie
+        for stream_offset, c in enumerate(iterator()):
+            n = state
+
+            while c not in n and n is not self.trie:
+                n = n.fall
+
+            if n is self.trie:
+                if c in n:
+                    n = n[c]
+            else:
+                n = n[c]
+
+            state = n
+
+            while n is not self.trie:
+                yield from ((stream_offset - len(source) + 1, source) for source in n.sources)
+                n = n.fall
+
+
+class StartsWithMatcher:
+    def __init__(self, *sequences_to_find):
+        self.trie = TrieNode()
+        for seq in sequences_to_find:
+            self.trie.add(seq)
+
+    def search(self, source_sequence: Union[Sequence, IO]):
+        if hasattr(source_sequence, 'read'):
+            def iterator():
+                while True:
+                    b = source_sequence.read(1)
+                    if not b:
+                        return
+                    yield b[0]
+        else:
+            def iterator():
+                return iter(source_sequence)
+
+        state = self.trie
+        yield from ((0, s) for s in state.sources)
+
+        for c in iterator():
+            if c not in state:
+                return
+
+            state = state[c]
+            yield from ((0, s) for s in state.sources)
+
+
+if __name__ == '__main__':
+    root = TrieNode()
+    root.add('The quick brown fox jumps over the lazy dog')
+    root.add('The quick person')
+    root.add('The best')
+    assert len(list(root.find_prefix('The'))) == 3
+    assert len(list(root.find_prefix('The quick'))) == 2
+    assert not root.find('The')
+    assert root.find('The best')
+
+    mss = MultiSequenceSearch(b'hack', b'hacker', b'crack', b'ack', b'kool')
+    to_search = b'This is a test to see if hack or hacker is in this string.'\
+                b'Can you crack it? If so, please ack, \'cause that would be kool.'
+    for offset, match in mss.search(to_search):
+        print(offset, match)
+        assert to_search[offset:offset+len(match)] == match
+
+    swm = StartsWithMatcher(b'hack', b'hacker', b'crack', b'ack', b'kool')
+    for match in swm.search(b'hacker'):
+        print(match)
+
+    print(ACNode.load(mss.trie.serialize()))

polyfile/serialization.py

@@ -0,0 +1,323 @@
+from enum import Enum
+import io
+from itertools import chain
+
+
+def write_int(i:int, stream):
+    bits = max(i.bit_length(), 1)
+    for bit_group in reversed(range((bits - 1) // 7 + 1)):
+        b = (i >> (bit_group * 7)) & 0b01111111
+        if bit_group == 0:
+            b |= 0b10000000
+        stream.write(bytes([b]))
+
+
+def read_int(stream):
+    i = 0
+    while True:
+        i <<= 7
+        b = ord(stream.read(1))
+        i |= (b & 0b01111111)
+        if b & 0b10000000:
+            break
+    return i
+
+ENCODINGS_BY_TYPE = {}
+ENCODINGS_BY_ID = {}
+CUSTOM_ENCODINGS = {}
+CUSTOM_ENCODINGS_BY_ID = {}
+
+
+def encode_int(i, _, stream):
+    write_int(i, stream)
+
+
+def decode_int(_, stream):
+    return read_int(stream)
+
+
+def encode_bytes(s, _, stream):
+    write_int(len(s), stream)
+    stream.write(s)
+
+
+def encode_string(s:str, _, stream):
+    encode_bytes(s.encode('utf-8'), None, stream)
+
+
+def decode_bytes(_, stream):
+    length = read_int(stream)
+    return stream.read(length)
+
+
+def decode_string(_, stream):
+    return decode_bytes(None, stream).decode('utf-8')
+
+
+def encode_list(lst, obj_ids, stream):
+    for obj in lst:
+        write_int(obj_ids[id(obj)], stream)
+    write_int(EncodeTypes.END.encoding_id, stream)
+
+
+encode_tuple = encode_list
+
+
+encode_set = encode_list
+
+
+encode_frozenset = encode_list
+
+
+def _decode_to_end(stream):
+    while True:
+        i = read_int(stream)
+        if i == EncodeTypes.END.encoding_id:
+            break
+        yield i
+
+
+def pairwise(iterable):
+    while True:
+        try:
+            a, b = next(iterable), next(iterable)
+        except StopIteration:
+            break
+        yield a, b
+
+
+def decode_list(objs, stream):
+    return [objs[i] for i in _decode_to_end(stream)]
+
+
+def decode_tuple(*args, **kwargs):
+    return tuple(decode_list(*args, **kwargs))
+
+
+def decode_set(*args, **kwargs):
+    return set(decode_list(*args, **kwargs))
+
+
+def decode_frozenset(*args, **kwargs):
+    return frozenset(decode_list(*args, **kwargs))
+
+
+def encode_dict(d: dict, obj_ids, stream):
+    for k, v in d.items():
+        write_int(obj_ids[id(k)], stream)
+        write_int(obj_ids[id(v)], stream)
+    write_int(EncodeTypes.END.encoding_id, stream)
+
+
+def decode_dict(objs, stream):
+    return {objs[k]: objs[v] for k, v in pairwise(_decode_to_end(stream))}
+
+
+def encode_none(*args, **kwargs):
+    pass
+
+
+def decode_none(*args, **kwargs):
+    return None
+
+
+def encode_bool(b: bool, _, stream):
+    write_int(int(b), stream)
+
+
+def decode_bool(_, stream):
+    return bool(read_int(stream))
+
+
+def encode_custom(c, obj_ids, stream):
+    write_int(CUSTOM_ENCODINGS[type(c)], stream)
+    encode_list(c.serialize(), obj_ids, stream)
+
+
+def decode_custom(objs, stream):
+    custom_type = CUSTOM_ENCODINGS_BY_ID[read_int(stream)]
+    args = decode_list(objs, stream)
+    return custom_type(*args)
+
+
+class _EndObject:
+    pass
+
+
+class _CustomObject:
+    pass
+
+
+def serializable(SerializableClass):
+    if not hasattr(SerializableClass, 'serialize'):
+        raise ValueError(f'Serializable class {SerializableClass} must implement the `serialize` member function')
+    new_id = len(CUSTOM_ENCODINGS)
+    CUSTOM_ENCODINGS[SerializableClass] = new_id
+    CUSTOM_ENCODINGS_BY_ID[new_id] = SerializableClass
+    return SerializableClass
+
+
+class EncodeTypes(Enum):
+    END = (0, _EndObject, None, None, None)
+    NONE = (1, type(None), encode_none, decode_none, lambda _: ())
+    BOOL = (2, bool, encode_bool, decode_bool, lambda _: ())
+    INT = (3, int, encode_int, decode_int, lambda _: ())
+    STRING = (4, str, encode_string, decode_string, lambda _: ())
+    BYTES = (5, bytes, encode_bytes, decode_bytes, lambda _: ())
+    LIST = (6, list, encode_list, decode_list, iter)
+    TUPLE = (7, tuple, encode_tuple, decode_tuple, iter)
+    DICT = (8, dict, encode_dict, decode_dict, lambda d: chain(d.keys(), d.values()))
+    SET = (9, set, encode_set, decode_set, iter)
+    FROZENSET = (10, frozenset, encode_frozenset, decode_frozenset, iter)
+    CUSTOM = (11, _CustomObject, encode_custom, decode_custom, lambda c: c.serialize())
+
+    def __init__(self, encoding_id, source_type, encoding_function, decoding_function, children):
+        self.encoding_id = encoding_id
+        self.source_type = source_type
+        self.encode = encoding_function
+        self.decode = decoding_function
+        self.children = children
+        ENCODINGS_BY_TYPE[self.source_type] = self
+        ENCODINGS_BY_ID[self.encoding_id] = self
+
+    @staticmethod
+    def get_by_type(source_type):
+        if source_type not in ENCODINGS_BY_TYPE:
+            return None
+        return ENCODINGS_BY_TYPE[source_type]
+
+    @staticmethod
+    def get_by_id(encoding_id):
+        if encoding_id not in ENCODINGS_BY_ID:
+            return None
+        return ENCODINGS_BY_ID[encoding_id]
+
+
+class EncodingError(RuntimeError):
+    def __init__(self, *args):
+        super().__init__(*args)
+
+
+def encode(obj, stream):
+    stack = [obj]
+    # Object ID 0 is reserved
+    obj_ids = {
+        0: None
+    }
+    objs = [
+        (None, None)
+    ]
+
+    # Do a DFS traversal to assign IDs
+    while stack:
+        s = stack.pop()
+        encoding = EncodeTypes.get_by_type(type(s))
+        if encoding is None:
+            # See if there is a custom encoding:
+            if type(s) in CUSTOM_ENCODINGS:
+                encoding = EncodeTypes.CUSTOM
+            else:
+                raise EncodingError(f"No encoding implemented for objects of type {type(s)}")
+        already_expanded = id(s) in obj_ids
+        if already_expanded:
+            old_id = obj_ids[id(s)]
+        # We add one to new_id here because obj_id zero is reserved
+        new_id = len(objs)
+        obj_ids[id(s)] = new_id
+        objs.append((s, encoding))
+        if already_expanded:
+            objs[old_id] = (None, None)
+            # Remove all of the dependencies, since they now need to be added after:
+            # TODO: eventually do cycle detection here, and throw an error instead of livelock
+            for child in encoding.children(s):
+                if id(child) in obj_ids:
+                    objs[obj_ids[id(child)]] = (None, None)
+                    del obj_ids[id(child)]
+        stack.extend(encoding.children(s))
+
+    # write the objects to the stream in reverse
+    n = len(objs)
+    for i, (s, encoding) in enumerate(reversed(objs)):
+        obj_id = n - i - 1
+        if encoding is None:
+            # This object was referenced twice, and this was a duplicate
+            continue
+        write_int(obj_id, stream)
+        write_int(encoding.encoding_id, stream)
+        encoding.encode(s, obj_ids, stream)
+
+
+def decode(stream):
+    objs = {}
+    while True:
+        b = stream.read(1)
+        if b is None or not b:
+            break
+        stream.seek(-1, 1)
+        obj_id = read_int(stream)
+        assert obj_id not in objs
+        encoding_id = read_int(stream)
+        encoding = EncodeTypes.get_by_id(encoding_id)
+        if encoding is None:
+            raise EncodingError(f"Unexpected encoded object of type #{obj_id}")
+        objs[obj_id] = encoding.decode(objs, stream)
+    return objs[1]
+
+
+def dump(obj, stream):
+    return encode(obj, stream)
+
+
+def load(stream):
+    return decode(stream)
+
+
+def dumps(obj):
+    stream = io.BytesIO()
+    dump(obj, stream)
+    return stream.getvalue()
+
+
+def loads(string:bytes):
+    with io.BytesIO(string) as stream:
+        return load(stream)
+
+
+if __name__ == '__main__':
+    ref = ["list", "used", "twice"]
+    test_obj = {
+        'testing': {
+            'foo': {10},
+            'bar': [1, 2, 3, b'1234\xFF', True, False],
+            'ref': ref
+        },
+        'baz': [
+            'a', ('b',), 'c',
+            {'d': 5},
+            None,
+            frozenset([1, 1, 2, 3, 5, 8])
+        ],
+        'ref': ref
+    }
+
+    #print(test_obj)
+    #encoded = encode(test_obj)
+    #print(encoded)
+    #decoded = decode(encoded)
+    #print(decoded)
+
+    for i in range(0, 2048):
+        assert loads(dumps(i)) == i
+
+    b = b"The quick brown fox jumps over the lazy dog"
+    s = b.decode('utf-8')
+    assert loads(dumps(b)) == b
+    assert loads(dumps(s)) == s
+
+    lst = list(range(1024))
+    assert loads(dumps(lst)) == lst
+
+    encoded = dumps(test_obj)
+    print(f"Encoded bytes: {len(encoded)}")
+
+    print(loads(encoded))

polyfile/structmatcher.py

@@ -0,0 +1,46 @@
+from typing import Iterator, Optional
+
+from .fileutils import Tempfile
+from .polyfile import InvalidMatch, Match, Matcher, Submatch
+from .structs import Field, Struct, StructError
+
+
+class PolyFileStruct(Struct):
+    __name__: Optional[str] = None
+
+    @property
+    def match_name(self) -> str:
+        if self.__name__ is not None:
+            return self.__name__
+        else:
+            return self.__class__.__name__
+
+    def match(self, matcher: Matcher, parent: Optional[Match] = None) -> Iterator[Submatch]:
+        m = Submatch(
+            self.match_name,
+            match_obj=self,
+            relative_offset=self.start_offset,
+            length=self.num_bytes,
+            parent=parent,
+            matcher=matcher
+        )
+        yield m
+        for field_name in self.fields.keys():
+            value: Field = getattr(self, field_name)
+            s = Submatch(
+                field_name,
+                match_obj=value,
+                relative_offset=value.start_offset - self.start_offset,
+                length=value.num_bytes,
+                parent=m,
+                matcher=matcher
+            )
+            yield s
+            try:
+                if isinstance(value, PolyFileStruct):
+                    yield from value.match(matcher, s)
+                elif isinstance(value, bytes):
+                    with Tempfile(value) as tmp:
+                        yield from matcher.match(tmp, parent=s)
+            except (InvalidMatch, StructError):
+                pass