<commit_before>__all__ = ['Table']

import os
import sys
from pprint import pprint
from collections import defaultdict

from .column import Column
from .schema import Schema
from .memtable import MemTable
from .sstable import SSTable
from .query import Query
from .deferred import Deferred
from .expr import Expr

class Table(object):
    MEMTABLE_LIMIT_N_ITEMS = 100

    def __init__(self, db, table_name, type_fields=None):
        self.store = db.store
        self.db = db
        self.table_name = table_name
        self.opened = False

        # schema
        self.schema = Schema(self, type_fields)

        # memtable
        self.memtable = MemTable(self)

        # sstables
        self.sstables = []
        table_path = self.get_path()

        for filename in os.listdir(table_path):
            if not filename.startswith('sstable-'):
                continue

            s = filename.index('sstable-') + len('sstable-')
            e = filename.index('.data')
            t = filename[s:e]

            # sstable
            sst = SSTable(self, t)
            sst.open() # FIXME: lazy open in SSTable only if required
            self.sstables.append(sst)

    def __getattr__(self, attr):
        c = getattr(self.schema, attr)
        return c

    def get_path(self):
        return os.path.join(self.db.get_path(), self.table_name)

    def is_opened(self):
        return self.opened

    def open(self):
        self.opened = True

    def close(self):
        for sst in self.sstables:
            if sst.is_opened():
                sst.close()

        self.opened = False

    def commit_if_required(self):
        if len(self.memtable) >= self.MEMTABLE_LIMIT_N_ITEMS:
            self.commit()

    def commit(self):
        # get sorted rows by primary_key
        columns = self.schema.primary_key
        rows = self.memtable.get_sorted_rows(columns)

        # create new sstable
        sst = SSTable(self, rows=rows)
        sst.open()
        self.sstables.append(sst)

        # clear memtable
        self.memtable = MemTable(self)

    @property
    def query(self):
        q = Query(self.db.store)
        return q

    def insert(self, **row):
        # tx
        tx = self.store.get_current_transaction()
        tx.log((self.db, self, Table._commit_insert, (self,), row))

    def _commit_insert(self, **row):
        # check if all columns exist in table's schema
        # compare against schema
        for k, v in row.items():
            if k not in self.schema:
                raise Exception(
                    'filed %r is not defined in schema for table %r' % (
                        k,
                        self.table_name,
                    )
                )

        # set default columns to None
        for k, v in self.schema:
            if k not in row:
                row[k] = None

        # build key
        key = tuple(row[k] for k in self.schema.primary_key)

        # insert key
        self.memtable.set(key, row)

        # commit if required
        self.commit_if_required()

    def get(self, *args):
        # key
        key = args

        # deferred
        d = Deferred()

        # tx
        tx = self.store.get_current_transaction()
        tx.log((self.db, self.table, Table._commit_get, (self, d, key), {}))

        return d

    def _commit_get(self, d, key):
        v, op, sp = self._get(key)
        d.set(v)

    def select(self, *args):
        # deferred, queue
        d = Deferred()
        q = Query(self.store, d)

        # tx
        tx = self.store.get_current_transaction()
        tx.log((self.db, self.table, Table._commit_select, (self, d, q), {}))

        return q

    def _cmp_ranges(self, a, b):
        ac, aop, av = a
        bc, bop, bv = b

        if ac == bc:
            if aop in ('<', '<=') and bop in ('<', '<='):
                return cmp(av, bv)
            elif aop in ('>=', '>') and bop in ('>', '>='):
                return cmp(av, bv)
            elif aop in ('>=', '>') and bop in ('<', '<='):
                if av < bv:
                    return -1
                elif av > bv:
                    return 1
                else:
                    return 0
            elif aop in ('<', '<=') and bop in ('>=', '>'):
                if av < bv:
                    return 0
                elif av > bv:
                    return 1
                else:
                    return 0
            else:
                return 0
        else:
            return cmp(ac, bc)

    def _eval_expr(self, expr):
        print '_eval_expr:', expr
        ranges = []

        if expr.op == 'and':
            l_ranges = self._eval_expr(expr.left)
            r_ranges = self._eval_expr(expr.right)
            _ranges = list(set(l_ranges + r_ranges))
            _ranges.sort(cmp=lambda a, b: self._cmp_ranges(a, b))

            # for same column, if op is '==',
            # then it cancels all other ops if found inside of a range
            ranges_by_column = {}

            for r in _ranges:
                c, op, v = r

                if c in ranges_by_column:
                    rs = ranges_by_column[c]
                    rs.append(r)
                else:
                    ranges_by_column[c] = [r]

            print 'ranges_by_column[0]:', ranges_by_column

            # optimize ranges by column
            for c, rs in ranges_by_column.items():
                if len(rs) == 1:
                    continue

                _rs = []

                for a, b in zip(rs[:-1], rs[1:]):
                    ac, aop, av = a
                    bc, bop, bv = b
                    print 'a, b:', a, b

                    if aop == '<':
                        if bop == '<':
                            r = (ac, '<', min(av, bv))
                            _rs.append(r)
                        elif bop == '<=':
                            if av <= bv:
                                _rs.append(a)
                            else:
                                _rs.append(b)
                        elif bop == '==':
                            if av > bv:
                                _rs.append(b)
                        elif bop == '>=':
                            if av > bv:
                                _rs.append(b)
                                _rs.append(a)
                        elif bop == '>':
                            if av > bv:
                                _rs.append(b)
                                _rs.append(a)
                    elif aop == '<=':
                        if bop == '<':
                            if av < bv:
                                _rs.append(a)
                            else:
                                _rs.append(b)
                        elif bop == '<=':
                            r = (ac, '<=', min(av, bv))
                            _rs.append(r)
                        elif bop == '==':
                            if av >= bv:
                                _rs.append(b)
                        elif bop == '>=':
                            if av == bv:
                                r = (ac, '==', av)
                                _rs.append(r)
                            elif av > bv:
                                _rs.append(b)
                                _rs.append(a)
                        elif bop == '>':
                            if av > bv:
                                _rs.append(b)
                                _rs.append(a)
                    elif aop == '==':
                        if bop == '<':
                            if av < bv:
                                _rs.append(a)
                        elif bop == '<=':
                            if av <= bv:
                                _rs.append(a)
                        elif bop == '==':
                            if av == bv:
                                _rs.append(a)
                        elif bop == '>=':
                            if av >= bv:
                                _rs.append(a)
                        elif bop == '>':
                            if av > bv:
                                _rs.append(a)
                    elif aop == '>=':
                        if bop == '<':
                            if av < bv:
                                _rs.append(a)
                                _rs.append(b)
                        elif bop == '<=':
                            if av == bv:
                                r = (ac, '==', av)
                                _rs.append(r)
                            elif av < bv:
                                _rs.append(a)
                                _rs.append(b)
                        elif bop == '==':
                            if av == bv:
                                _rs.append(b)
                        elif bop == '>=':
                            r = (ac, '>=', max(av, bv))
                            _rs.append(r)
                        elif bop == '>':
                            if av <= bv:
                                _rs.append(b)
                            else:
                                _rs.append(a)
                    elif aop == '>':
                        if bop == '<':
                            if av < bv:
                                _rs.append(a)
                                _rs.append(b)
                        elif bop == '<=':
                            if av < bv:
                                _rs.append(a)
                                _rs.append(b)
                        elif bop == '==':
                            if av < bv:
                                _rs.append(b)
                        elif bop == '>=':
                            if av < bv:
                                _rs.append(b)
                            else:
                                _rs.append(a)
                        elif bop == '>':
                            r = (ac, '>=', max(av, bv))
                            _rs.append(r)

                ranges_by_column[c] = _rs

            print 'ranges_by_column[1]:', ranges_by_column
            print '_ranges:', _ranges
            ranges.extend(_ranges)
        elif expr.op == 'or':
            # FIXME: OR is not planned for now, but have it in mind
            pass
        else:
            k = (expr.right,)
            cs = (expr.left.name,)

            if expr.op == '==':
                v, op, sp = self._get(k, cs)
            elif expr.op == '<':
                v, op, sp = self._get_lt(k, cs)
            elif expr.op == '<=':
                v, op, sp = self._get_le(k, cs)
            elif expr.op == '>':
                v, op, sp = self._get_gt(k, cs)
            elif expr.op == '>=':
                v, op, sp = self._get_ge(k, cs)

            ranges.append((cs[0], expr.op, k[0]))

        return ranges

    def _commit_select(self, d, q):
        rows = []   # get from keys
        keys = []   # get from ranges
        ranges = self._eval_expr(q.where_clause)
        print 'ranges:', ranges
        d.set(rows)

    def _get(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp

    def _get_lt(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get_lt(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get_lt(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp

    def _get_le(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get_le(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get_le(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp

    def _get_gt(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get_gt(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get_gt(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp

    def _get_ge(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get_ge(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get_ge(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp
<commit_msg>store: Table, optimized ranges<commit_after>__all__ = ['Table']

import os
import sys
from pprint import pprint
from collections import defaultdict

from .column import Column
from .schema import Schema
from .memtable import MemTable
from .sstable import SSTable
from .query import Query
from .deferred import Deferred
from .expr import Expr

class Table(object):
    MEMTABLE_LIMIT_N_ITEMS = 100

    def __init__(self, db, table_name, type_fields=None):
        self.store = db.store
        self.db = db
        self.table_name = table_name
        self.opened = False

        # schema
        self.schema = Schema(self, type_fields)

        # memtable
        self.memtable = MemTable(self)

        # sstables
        self.sstables = []
        table_path = self.get_path()

        for filename in os.listdir(table_path):
            if not filename.startswith('sstable-'):
                continue

            s = filename.index('sstable-') + len('sstable-')
            e = filename.index('.data')
            t = filename[s:e]

            # sstable
            sst = SSTable(self, t)
            sst.open() # FIXME: lazy open in SSTable only if required
            self.sstables.append(sst)

    def __getattr__(self, attr):
        c = getattr(self.schema, attr)
        return c

    def get_path(self):
        return os.path.join(self.db.get_path(), self.table_name)

    def is_opened(self):
        return self.opened

    def open(self):
        self.opened = True

    def close(self):
        for sst in self.sstables:
            if sst.is_opened():
                sst.close()

        self.opened = False

    def commit_if_required(self):
        if len(self.memtable) >= self.MEMTABLE_LIMIT_N_ITEMS:
            self.commit()

    def commit(self):
        # get sorted rows by primary_key
        columns = self.schema.primary_key
        rows = self.memtable.get_sorted_rows(columns)

        # create new sstable
        sst = SSTable(self, rows=rows)
        sst.open()
        self.sstables.append(sst)

        # clear memtable
        self.memtable = MemTable(self)

    @property
    def query(self):
        q = Query(self.db.store)
        return q

    def insert(self, **row):
        # tx
        tx = self.store.get_current_transaction()
        tx.log((self.db, self, Table._commit_insert, (self,), row))

    def _commit_insert(self, **row):
        # check if all columns exist in table's schema
        # compare against schema
        for k, v in row.items():
            if k not in self.schema:
                raise Exception(
                    'filed %r is not defined in schema for table %r' % (
                        k,
                        self.table_name,
                    )
                )

        # set default columns to None
        for k, v in self.schema:
            if k not in row:
                row[k] = None

        # build key
        key = tuple(row[k] for k in self.schema.primary_key)

        # insert key
        self.memtable.set(key, row)

        # commit if required
        self.commit_if_required()

    def get(self, *args):
        # key
        key = args

        # deferred
        d = Deferred()

        # tx
        tx = self.store.get_current_transaction()
        tx.log((self.db, self.table, Table._commit_get, (self, d, key), {}))

        return d

    def _commit_get(self, d, key):
        v, op, sp = self._get(key)
        d.set(v)

    def select(self, *args):
        # deferred, queue
        d = Deferred()
        q = Query(self.store, d)

        # tx
        tx = self.store.get_current_transaction()
        tx.log((self.db, self.table, Table._commit_select, (self, d, q), {}))

        return q

    def _cmp_ranges(self, a, b):
        ac, aop, av = a
        bc, bop, bv = b

        if ac == bc:
            if aop in ('<', '<=') and bop in ('<', '<='):
                return cmp(av, bv)
            elif aop in ('>=', '>') and bop in ('>', '>='):
                return cmp(av, bv)
            elif aop in ('>=', '>') and bop in ('<', '<='):
                if av < bv:
                    return -1
                elif av > bv:
                    return 1
                else:
                    return 0
            elif aop in ('<', '<=') and bop in ('>=', '>'):
                if av < bv:
                    return 0
                elif av > bv:
                    return 1
                else:
                    return 0
            else:
                return 0
        else:
            return cmp(ac, bc)

    def _eval_expr(self, expr):
        print '_eval_expr:', expr
        ranges = []

        if expr.op == 'and':
            l_ranges = self._eval_expr(expr.left)
            r_ranges = self._eval_expr(expr.right)
            _ranges = list(set(l_ranges + r_ranges))
            _ranges.sort(cmp=lambda a, b: self._cmp_ranges(a, b))
            print '_ranges[0]:', _ranges

            # for same column, if op is '==',
            # then it cancels all other ops if found inside of a range
            ranges_by_column = {}

            for r in _ranges:
                c, op, v = r

                if c in ranges_by_column:
                    rs = ranges_by_column[c]
                    rs.append(r)
                else:
                    ranges_by_column[c] = [r]

            print 'ranges_by_column[0]:', ranges_by_column

            # optimize ranges by column
            for c, rs in ranges_by_column.items():
                if len(rs) == 1:
                    continue

                _rs = []

                for a, b in zip(rs[:-1], rs[1:]):
                    ac, aop, av = a
                    bc, bop, bv = b
                    print 'a, b:', a, b

                    if aop == '<':
                        if bop == '<':
                            r = (ac, '<', min(av, bv))
                            _rs.append(r)
                        elif bop == '<=':
                            if av <= bv:
                                _rs.append(a)
                            else:
                                _rs.append(b)
                        elif bop == '==':
                            if av > bv:
                                _rs.append(b)
                        elif bop == '>=':
                            if av > bv:
                                _rs.append(b)
                                _rs.append(a)
                        elif bop == '>':
                            if av > bv:
                                _rs.append(b)
                                _rs.append(a)
                    elif aop == '<=':
                        if bop == '<':
                            if av < bv:
                                _rs.append(a)
                            else:
                                _rs.append(b)
                        elif bop == '<=':
                            r = (ac, '<=', min(av, bv))
                            _rs.append(r)
                        elif bop == '==':
                            if av >= bv:
                                _rs.append(b)
                        elif bop == '>=':
                            if av == bv:
                                r = (ac, '==', av)
                                _rs.append(r)
                            elif av > bv:
                                _rs.append(b)
                                _rs.append(a)
                        elif bop == '>':
                            if av > bv:
                                _rs.append(b)
                                _rs.append(a)
                    elif aop == '==':
                        if bop == '<':
                            if av < bv:
                                _rs.append(a)
                        elif bop == '<=':
                            if av <= bv:
                                _rs.append(a)
                        elif bop == '==':
                            if av == bv:
                                _rs.append(a)
                        elif bop == '>=':
                            if av >= bv:
                                _rs.append(a)
                        elif bop == '>':
                            if av > bv:
                                _rs.append(a)
                    elif aop == '>=':
                        if bop == '<':
                            if av < bv:
                                _rs.append(a)
                                _rs.append(b)
                        elif bop == '<=':
                            if av == bv:
                                r = (ac, '==', av)
                                _rs.append(r)
                            elif av < bv:
                                _rs.append(a)
                                _rs.append(b)
                        elif bop == '==':
                            if av == bv:
                                _rs.append(b)
                        elif bop == '>=':
                            r = (ac, '>=', max(av, bv))
                            _rs.append(r)
                        elif bop == '>':
                            if av <= bv:
                                _rs.append(b)
                            else:
                                _rs.append(a)
                    elif aop == '>':
                        if bop == '<':
                            if av < bv:
                                _rs.append(a)
                                _rs.append(b)
                        elif bop == '<=':
                            if av < bv:
                                _rs.append(a)
                                _rs.append(b)
                        elif bop == '==':
                            if av < bv:
                                _rs.append(b)
                        elif bop == '>=':
                            if av < bv:
                                _rs.append(b)
                            else:
                                _rs.append(a)
                        elif bop == '>':
                            r = (ac, '>=', max(av, bv))
                            _rs.append(r)

                ranges_by_column[c] = _rs

            print 'ranges_by_column[1]:', ranges_by_column

            _ranges = []

            for c in sorted(ranges_by_column.keys()):
                rs = ranges_by_column[c]
                _ranges.extend(rs)

            print '_ranges[1]:', _ranges
            ranges.extend(_ranges)
        elif expr.op == 'or':
            # FIXME: OR is not planned for now, but have it in mind
            pass
        else:
            k = (expr.right,)
            cs = (expr.left.name,)

            if expr.op == '==':
                v, op, sp = self._get(k, cs)
            elif expr.op == '<':
                v, op, sp = self._get_lt(k, cs)
            elif expr.op == '<=':
                v, op, sp = self._get_le(k, cs)
            elif expr.op == '>':
                v, op, sp = self._get_gt(k, cs)
            elif expr.op == '>=':
                v, op, sp = self._get_ge(k, cs)

            ranges.append((cs[0], expr.op, k[0]))

        return ranges

    def _commit_select(self, d, q):
        rows = []   # get from keys
        keys = []   # get from ranges
        ranges = self._eval_expr(q.where_clause)
        print 'ranges:', ranges
        d.set(rows)

    def _get(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp

    def _get_lt(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get_lt(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get_lt(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp

    def _get_le(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get_le(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get_le(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp

    def _get_gt(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get_gt(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get_gt(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp

    def _get_ge(self, key, columns=None):
        try:
            v, op, sp = self.memtable.get_ge(key, columns)
        except KeyError as e:
            for sst in reversed(self.sstables):
                try:
                    v, op, sp = sst.get_ge(key, columns)
                    break
                except KeyError as e:
                    pass
            else:
                raise KeyError

        return v, op, sp
<|endoftext|>


<commit_before>import rethinkdb as r

from adapt.engine import IntentDeterminationEngine
from adapt.intent import IntentBuilder

def register_intent(name, engine, *keywords, **kwargs):

    for keyword in keywords:
        engine.register_entity(keyword, name.title() + "Keyword")

    for index, values in kwargs.items():
        for value in values:
            engine.register_entity(value, name.title() + index.title())

    intent = IntentBuilder(name.title() + "Intent")\
        .require("RecycleKeyword")\
        .require(name.title() + "Keyword")

    for index in kwargs:
        intent = intent.optionally(name.title() + index.title())

    intent = intent.build()

    engine.register_intent_parser(intent)

def build_engine(rdb_conn):
    """Build a recycling intent determination engine."""

    engine = IntentDeterminationEngine()

    recycle_keywords = ["recycle", "recyclable"]

    for rk in recycle_keywords:
        engine.register_entity(rk, "RecycleKeyword")

    plastic_keywords = [item["name"] for item in r.table(
        "items").filter({"type": "plastic"}).run(rdb_conn)]

    register_intent(
        "plastic", engine,
        *plastic_keywords,
        descriptor=["number"],
        numbers=["1", "2", "3", "4", "5", "6", "7"]
    )

    glass_keywords = [item["name"] for item in r.table(
        "items").filter({"type": "glass"}).run(rdb_conn)]

    register_intent(
        "glass", engine,
        *glass_keywords
    )

    paper_keywords = [item["name"] for item in r.table(
        "items").filter({"type": "paper"}).run(rdb_conn)]

    register_intent(
        "paper", engine,
        *paper_keywords
    )

    return engine

def parse(input_text, *, engine):
    """
    Takes in input text and parses it to determine whether it is talking
    about plastic, glass, or paper, and then to find any refining things
    like the number of plastic or the type of glass.
    """

    for intent in engine.determine_intent(input_text):
        if intent is not None and intent.get('confidence') > 0:
            yield intent

if __name__ == "__main__":

    tests = [
        "Can I recycle this glass bottle?",
        "Is a number 6 bottle recyclable?",
        "Can I recycle this box?"
    ]

    _engine = build_engine(r.connect(host="localhost", port=28015, db="swish"))

    for test in tests:
        print(test)
        print(list(parse(test, engine=_engine)))
        print()
<commit_msg>Add lemmatization<commit_after>import rethinkdb as r
from nltk.stem import WordNetLemmatizer

from adapt.engine import IntentDeterminationEngine
from adapt.intent import IntentBuilder

wnl = WordNetLemmatizer()

def register_intent(name, engine, *keywords, **kwargs):

    for keyword in keywords:
        engine.register_entity(keyword, name.title() + "Keyword")

    for index, values in kwargs.items():
        for value in values:
            engine.register_entity(value, name.title() + index.title())

    intent = IntentBuilder(name.title() + "Intent")\
        .require("RecycleKeyword")\
        .require(name.title() + "Keyword")

    for index in kwargs:
        intent = intent.optionally(name.title() + index.title())

    intent = intent.build()

    engine.register_intent_parser(intent)

def build_engine(rdb_conn):
    """Build a recycling intent determination engine."""

    engine = IntentDeterminationEngine()

    recycle_keywords = ["recycle", "recyclable"]

    for rk in recycle_keywords:
        engine.register_entity(rk, "RecycleKeyword")

    plastic_keywords = [item["name"] for item in r.table(
        "items").filter({"type": "plastic"}).run(rdb_conn)]

    register_intent(
        "plastic", engine,
        *plastic_keywords,
        descriptor=["number"],
        numbers=["1", "2", "3", "4", "5", "6", "7"]
    )

    glass_keywords = [item["name"] for item in r.table(
        "items").filter({"type": "glass"}).run(rdb_conn)]

    register_intent(
        "glass", engine,
        *glass_keywords
    )

    paper_keywords = [item["name"] for item in r.table(
        "items").filter({"type": "paper"}).run(rdb_conn)]

    register_intent(
        "paper", engine,
        *paper_keywords
    )

    return engine

def parse(input_text, *, engine):
    """
    Takes in input text and parses it to determine whether it is talking
    about plastic, glass, or paper, and then to find any refining things
    like the number of plastic or the type of glass.
    """

    lemmatized = ' '.join(wnl.lemmatize(word) for word in input_text.split())

    for intent in engine.determine_intent(input_text):
        if intent is not None and intent.get('confidence') > 0:
            yield intent

if __name__ == "__main__":

    tests = [
        "Can I recycle this glass bottle?",
        "Is a number 6 bottle recyclable?",
        "Can I recycle this box?"
    ]

    _engine = build_engine(r.connect(host="localhost", port=28015, db="swish"))

    for test in tests:
        print(test)
        print(list(parse(test, engine=_engine)))
        print()
<|endoftext|>


<commit_before>import os
import pudb
import bitarray
import copy
from collections import namedtuple

start_end_pair = namedtuple('start_end_pair', 'start end')

def build_dirs(files):
    '''
    Build necessary directories based on a list of file paths
    '''

    for i in files:
        if len(i['path']) > 1:
            addpath = os.path.join(*i['path'][:-1])
            if addpath and addpath not in os.listdir(os.getcwd()):
                os.makedirs(addpath)
                print 'just made path', addpath

def get_want_file_pos(file_list):
    '''
    Ask the user which files in file_list he or she is interested in.
    Return indices for the files inside file_list
    '''
    want_file_pos = []
    print '\nFiles contained:\n'
    for i in file_list:
        print(os.path.join(*i['path']))
    while 1:
        all_answer = raw_input('\nDo you want all these files? (y/n): ')
        if all_answer in ('y', 'n'):
            break
    if all_answer == 'y':
        # TODO -- Can have something simpler here when user wants everything
        want_file_pos = range(len(file_list))
        return want_file_pos
    if all_answer == 'n':
        for j, tfile in enumerate(file_list):
            while 1:
                file_answer = raw_input('Do you want {}? '
                                        '(y/n): '.format(os.path.join
                                                        (*tfile['path'])))

                if file_answer in ('y', 'n'):
                    break
            if file_answer == 'y':
                want_file_pos.append(j)
        print "Here are all the files you want:"
        for k in want_file_pos:
            print os.path.join(*file_list[k]['path'])
        return want_file_pos

def get_file_starts(file_list):
    '''
    Return the starting position (in bytes) of a list of files by
    iteratively summing their lengths
    '''
    starts = []
    total = 0
    for i in file_list:
        starts.append(total)
        total += i['length']
    print starts
    return starts

def get_write_file(byte_index=0, file_starts=[0], files=[], outfiles=[]):
    '''
    Retrieve the actual file that the current block of data should be
    written to.
    '''
    i = 1
    while i <= len(file_starts):
        start = file_starts[-i]
        if start <= byte_index:
            tfile = files[-i]
            break
        else:
            i += 1
    if not tfile:
        pudb.set_trace()
    j = 1
    while j <= len(outfiles):
        if outfiles[-j].name == os.path.join(*tfile['path']):
            return outfiles[-j]
        else:
            j += 1
    else:
        # pudb.set_trace()
        raise Exception("Shit isn't matching")

def get_file_start(byte_index=0, file_starts=[]):
    '''
    Find the starting position of the earliest file that I want to write to
    '''
    # Seems like I should be counting forward through these. Find the first
    # file whose starting position is <= index?
    i = 1
    while i <= len(file_starts) + 1:
        if byte_index >= file_starts[-i]:
            return file_starts[-i]
        else:
            i += 1

def get_heads_tails(want_file_pos=[], file_starts=[], num_pieces=0,
                    piece_length=0):
    heads_tails = []
    for i in want_file_pos:
        head_tail = get_head_tail(want_index=i, file_starts=file_starts,
                                  num_pieces=num_pieces,
                                  piece_length=piece_length)
        heads_tails.append(head_tail)
    return heads_tails

def get_head_tail(want_index=0, file_starts=[], num_pieces=0,
                  piece_length=0):

    # Find the byte value where the file starts
    byte_start = file_starts[want_index]

    # The firt piece we care about is at the point where the combined length
    # is *just* less than or equal to byte_start
    first_piece = byte_start // piece_length

    # We want it in a separate variable so we can iterate
    piece_pos = first_piece

    # Find if we want the last file in the torrent
    if want_index == len(file_starts) - 1:
        last_piece = num_pieces - 1

    # Otherwise we want a different piece
    elif want_index < len(file_starts) - 1:
        next_file_start = file_starts[want_index + 1]
        while piece_pos * piece_length < next_file_start:
            piece_pos += 1

        # We want the piece *before* the first one after the next file starts
        last_piece = piece_pos - 1

    return start_end_pair(start=first_piece, end=last_piece)

def build_bitfield(heads_and_tails=[], num_pieces=0):
    this_bitfield = bitarray.bitarray('0' * num_pieces)
    for i in heads_and_tails:
        for j in range(i.start, i.end + 1):
            this_bitfield[j] = True
    return this_bitfield

class switchboard(object):
    def __init__(self, dirname='', file_list=[], piece_length=0, num_pieces=0):
        self.dirname = dirname
        self.file_list = copy.deepcopy(file_list)
        self.piece_length = piece_length
        self.num_pieces = num_pieces
        self.file_starts = get_file_starts(self.file_list)
        self.want_file_pos = get_want_file_pos(self.file_list)
        self.outfiles = []
        self.byte_index = 0
        os.mkdir(self.dirname)
        os.chdir(os.path.join(os.getcwd(), self.dirname))
        want_files = [self.file_list[index] for index in self.want_file_pos]
        build_dirs(want_files)
        for i in self.want_file_pos:
            thisfile = open(os.path.join(*self.file_list[i]
                                         ['path']), 'w')
            self.outfiles.append(thisfile)
        heads_and_tails = get_heads_tails(want_file_pos=self.want_file_pos,
                                          file_starts=self.file_starts,
                                          num_pieces=self.num_pieces,
                                          piece_length=self.piece_length)
        self.bitfield = build_bitfield(heads_and_tails,
                                       num_pieces=self.num_pieces)

    def seek(self, index):
        '''
        Set how far to advance (bytewise) in file list
        '''
        self.byte_index = index

    def write(self, block):

        # file_start is the byte offset of the rightmost file whose
        # offset is less than index. It's the offset of the file that
        # the block starting at byte_index should begin writing to.
        file_start = get_file_start(byte_index=self.byte_index,
                                    file_starts=self.file_starts)

        # write_file is the actual file that we ought to be writing to.
        write_file = get_write_file(byte_index=self.byte_index,
                                    files=self.file_list,
                                    file_starts=self.file_starts,
                                    outfiles=self.outfiles)

        file_internal_index = self.byte_index - file_start
        write_file.seek(file_internal_index)
        file_end = (self.file_starts
                    [self.file_starts.index(file_start) + 1])

        bytes_writable = file_end - file_internal_index
        if bytes_writable < len(block):
            # pudb.set_trace()
            write_file.write(block[:bytes_writable])
            block = block[bytes_writable:]

            # The next index in the entire torrent
            j = self.file_starts.index(file_start) + 1

            # If we still want a higher index
            if j < self.want_file_pos[-1]:

                # Index of current file among all files in torrent
                current_file_index = self.file_starts.index(file_start)

                # Index of that index inside want_file_pos
                inner_want_index = self.want_file_pos.index(current_file_index)

                # The index of the next file we want among all files in torent
                next_file_index = self.want_file_pos[inner_want_index + 1]

                # The starting byte value of that file
                self.byte_index = self.file_starts[next_file_index]

                # If there's enough block to make it there
                if len(block) > self.byte_index - file_start:
                    block = block[self.byte_index - file_start]

                    # Start writing again
                    self.write(block)

            else:
                return

        else:
            write_file.write(block)
            block = None

    def mark_off(self, index):
        self.bitfield[index] = False

    @property
    def complete(self):
        if any(self.bitfield):
            return False
        else:
            return True

    def close(self):
        for i in self.outfiles:
            i.close()

# def test_heads_tails():
#     files = [{'length': 3}, {'length': 3}, {'length': 3}, {'length': 3},
#              {'length': 3}]
#     want_file_pos = [1, 3]
#     piece_length = 4
#     file_length = 3
#     file_starts = [i * file_length for i in range(len(files))]
#     num_pieces = 4
#     print 'case 1'
#     heads_tails = get_heads_tails(want_file_pos=want_file_pos,
#                                   file_starts=file_starts,
#                                   num_pieces=num_pieces,
#                                   piece_length=piece_length)
#     assert build_bitfield(heads_tails, num_pieces) == bitarray.bitarray('1110')

#     files = [{'length': 4}, {'length': 4}, {'length': 4}, {'length': 4},
#              {'length': 4}]
#     want_file_pos = [1]
#     piece_length = 3
#     file_length = 4
#     file_starts = [i * file_length for i in range(len(files))]
#     num_pieces = 4
#     heads_tails = get_heads_tails(want_file_pos=want_file_pos,
#                                   file_starts=file_starts,
#                                   num_pieces=num_pieces,
#                                   piece_length=piece_length)
#     assert build_bitfield(heads_tails, num_pieces) == bitarray.bitarray('0110')

# if __name__ == '__main__':
#     test_heads_tails()
<commit_msg>Finally (almost) have subset downloads working. Bug screws up last few files.<commit_after>import os
import pudb
import bitarray
import copy
from collections import namedtuple

start_end_pair = namedtuple('start_end_pair', 'start end')

def build_dirs(files):
    '''
    Build necessary directories based on a list of file paths
    '''

    for i in files:
        if len(i['path']) > 1:
            addpath = os.path.join(*i['path'][:-1])
            if addpath and addpath not in os.listdir(os.getcwd()):
                os.makedirs(addpath)
                print 'just made path', addpath

def get_want_file_pos(file_list):
    '''
    Ask the user which files in file_list he or she is interested in.
    Return indices for the files inside file_list
    '''
    want_file_pos = []
    print '\nFiles contained:\n'
    for i in file_list:
        print(os.path.join(*i['path']))
    while 1:
        all_answer = raw_input('\nDo you want all these files? (y/n): ')
        if all_answer in ('y', 'n'):
            break
    if all_answer == 'y':
        # TODO -- Can have something simpler here when user wants everything
        want_file_pos = range(len(file_list))
        return want_file_pos
    if all_answer == 'n':
        for j, tfile in enumerate(file_list):
            while 1:
                file_answer = raw_input('Do you want {}? '
                                        '(y/n): '.format(os.path.join
                                                        (*tfile['path'])))

                if file_answer in ('y', 'n'):
                    break
            if file_answer == 'y':
                want_file_pos.append(j)
        print "Here are all the files you want:"
        for k in want_file_pos:
            print os.path.join(*file_list[k]['path'])
        return want_file_pos

def get_file_starts(file_list):
    '''
    Return the starting position (in bytes) of a list of files by
    iteratively summing their lengths
    '''
    starts = []
    total = 0
    for i in file_list:
        starts.append(total)
        total += i['length']
    print starts
    return starts

def get_rightmost_file(byte_index=0, file_starts=[0], files=[]):

    '''
    Retrieve the highest-indexed file that starts at or before byte_index.
    '''
    i = 1
    while i <= len(file_starts):
        start = file_starts[-i]
        if start <= byte_index:
            return files[-i]
        else:
            i += 1
    else:
        raise Exception('byte_index lower than all file_starts')

def get_next_want_file(byte_index=0, want_file_pos=[],
                       file_starts=[0], files=[], block=''):
    while block:
        rightmost = get_rightmost_file(byte_index=byte_index,
                                       file_starts=file_starts,
                                       files=files)
        if files.index(rightmost) in want_file_pos:
            return os.path.join(*rightmost['path'])
        else:
            byte_index = byte_index + rightmost['length']
            if len(block) > rightmost['length']:
                block = block[rightmost['length']:]
            else:
                block = ''
    else:
        return None

def get_file_start(byte_index=0, file_starts=[]):
    '''
    Find the starting position of the earliest file that I want to write to
    '''
    # Seems like I should be counting forward through these. Find the first
    # file whose starting position is <= index?
    i = 1
    while i <= len(file_starts) + 1:
        if byte_index >= file_starts[-i]:
            return file_starts[-i]
        else:
            i += 1

def get_heads_tails(want_file_pos=[], file_starts=[], num_pieces=0,
                    piece_length=0):
    heads_tails = []
    for i in want_file_pos:
        head_tail = get_head_tail(want_index=i, file_starts=file_starts,
                                  num_pieces=num_pieces,
                                  piece_length=piece_length)
        heads_tails.append(head_tail)
    return heads_tails

def get_head_tail(want_index=0, file_starts=[], num_pieces=0,
                  piece_length=0):

    # Find the byte value where the file starts
    byte_start = file_starts[want_index]

    # The firt piece we care about is at the point where the combined length
    # is *just* less than or equal to byte_start
    first_piece = byte_start // piece_length

    # We want it in a separate variable so we can iterate
    piece_pos = first_piece

    # Find if we want the last file in the torrent
    if want_index == len(file_starts) - 1:
        last_piece = num_pieces - 1

    # Otherwise we want a different piece
    elif want_index < len(file_starts) - 1:
        next_file_start = file_starts[want_index + 1]
        while piece_pos * piece_length < next_file_start:
            piece_pos += 1

        # We want the piece *before* the first one after the next file starts
        last_piece = piece_pos - 1

    return start_end_pair(start=first_piece, end=last_piece)

def build_bitfield(heads_and_tails=[], num_pieces=0):
    this_bitfield = bitarray.bitarray('0' * num_pieces)
    for i in heads_and_tails:
        for j in range(i.start, i.end + 1):
            this_bitfield[j] = True
    return this_bitfield

class switchboard(object):
    def __init__(self, dirname='', file_list=[], piece_length=0, num_pieces=0):
        self.dirname = dirname
        self.file_list = copy.deepcopy(file_list)
        self.piece_length = piece_length
        self.num_pieces = num_pieces
        self.file_starts = get_file_starts(self.file_list)
        self.want_file_pos = get_want_file_pos(self.file_list)
        self.outfiles = []
        self.byte_index = 0
        os.mkdir(self.dirname)
        os.chdir(os.path.join(os.getcwd(), self.dirname))
        want_files = [self.file_list[index] for index in self.want_file_pos]
        build_dirs(want_files)
        for i in self.want_file_pos:
            thisfile = open(os.path.join(*self.file_list[i]
                                         ['path']), 'w')
            self.outfiles.append(thisfile)
        heads_and_tails = get_heads_tails(want_file_pos=self.want_file_pos,
                                          file_starts=self.file_starts,
                                          num_pieces=self.num_pieces,
                                          piece_length=self.piece_length)
        self.bitfield = build_bitfield(heads_and_tails,
                                       num_pieces=self.num_pieces)

    def seek(self, index):
        '''
        Set how far to advance (bytewise) in file list
        '''
        self.byte_index = index

    def write(self, block):
        # file_start is the byte offset of the rightmost file whose
        # offset is less than index. It's the offset of the file that
        # the block starting at byte_index should begin writing to.
        file_start = get_file_start(byte_index=self.byte_index,
                                    file_starts=self.file_starts)

        # write_path is the file path that we ought to be writing to.
        write_path = get_next_want_file(byte_index=self.byte_index,
                                        files=self.file_list,
                                        file_starts=self.file_starts,
                                        block=block,
                                        want_file_pos=self.want_file_pos)

        if not write_path:
            return

        i = 0
        while i < len(self.outfiles):
            if self.outfiles[i].name == write_path:
                write_file = self.outfiles[i]
                break
            else:
                i += 1
        else:
            pudb.set_trace()
            raise Exception('Nothing matches')

        file_internal_index = self.byte_index - file_start
        write_file.seek(file_internal_index)
        file_end = (self.file_starts
                    [self.file_starts.index(file_start) + 1])

        bytes_writable = file_end - file_internal_index
        if bytes_writable < len(block):
            write_file.write(block[:bytes_writable])
            block = block[bytes_writable:]

            # The next index in the entire torrent
            j = self.file_starts.index(file_start) + 1

            # If we still want a higher index
            if j < self.want_file_pos[-1]:

                # Index of current file among all files in torrent
                current_file_index = self.file_starts.index(file_start)

                # Index of that index inside want_file_pos
                inner_want_index = self.want_file_pos.index(current_file_index)

                # The index of the next file we want among all files in torent
                next_file_index = self.want_file_pos[inner_want_index + 1]

                # The starting byte value of that file
                self.byte_index = self.file_starts[next_file_index]

                # If there's enough block to make it there
                if len(block) > self.byte_index - file_start:
                    block = block[self.byte_index - file_start]

                    # Start writing again
                    self.write(block)

            else:
                return

        else:
            write_file.write(block)
            block = None

    def mark_off(self, index):
        self.bitfield[index] = False

    @property
    def complete(self):
        if any(self.bitfield):
            return False
        else:
            return True

    def close(self):
        for i in self.outfiles:
            i.close()

# def test_get_write_file():
#     file_starts = [0, 483023, 944949, 4157426, 18877975, 36667180, 51499115,
#                    70003240, 84992746, 93577236, 109611431, 127573045,
#                    139208428, 151130741, 166870702, 180550961, 188662949,
#                    197822552, 214302710, 225269360, 236530253, 242925892,
#                    254454279, 265898029, 273484016, 273485312]

#     files = [{'path': ['Content', '174-h.htm'], 'length': 483023},
#              {'path': ['Content', '174.txt'], 'length': 461926},
#              {'path': ['Content', 'pictureofdoriangray_00_wilde_64kb.mp3'], 'length': 3212477},
#              {'path': ['Content', 'pictureofdoriangray_01_wilde_64kb.mp3'], 'length': 14720549},
#              {'path': ['Content', 'pictureofdoriangray_02_wilde_64kb.mp3'], 'length': 17789205},
#              {'path': ['Content', 'pictureofdoriangray_03_wilde_64kb.mp3'], 'length': 14831935},
#              {'path': ['Content', 'pictureofdoriangray_04_wilde_64kb.mp3'], 'length': 18504125},
#              {'path': ['Content', 'pictureofdoriangray_05_wilde_64kb.mp3'], 'length': 14989506},
#              {'path': ['Content', 'pictureofdoriangray_06_wilde_64kb.mp3'], 'length': 8584490},
#              {'path': ['Content', 'pictureofdoriangray_07_wilde_64kb.mp3'], 'length': 16034195},
#              {'path': ['Content', 'pictureofdoriangray_08_wilde_64kb.mp3'], 'length': 17961614},
#              {'path': ['Content', 'pictureofdoriangray_09_wilde_64kb.mp3'], 'length': 11635383},
#              {'path': ['Content', 'pictureofdoriangray_10_wilde_64kb.mp3'], 'length': 11922313},
#              {'path': ['Content', 'pictureofdoriangray_11a_wilde_64kb.mp3'], 'length': 15739961},
#              {'path': ['Content', 'pictureofdoriangray_11b_wilde_64kb.mp3'], 'length': 13680259},
#              {'path': ['Content', 'pictureofdoriangray_12_wilde_64kb.mp3'], 'length': 8111988},
#              {'path': ['Content', 'pictureofdoriangray_13_wilde_64kb.mp3'], 'length': 9159603},
#              {'path': ['Content', 'pictureofdoriangray_14_wilde_64kb.mp3'], 'length': 16480158},
#              {'path': ['Content', 'pictureofdoriangray_15_wilde_64kb.mp3'], 'length': 10966650},
#              {'path': ['Content', 'pictureofdoriangray_16_wilde_64kb.mp3'], 'length': 11260893},
#              {'path': ['Content', 'pictureofdoriangray_17_wilde_64kb.mp3'], 'length': 6395639},
#              {'path': ['Content', 'pictureofdoriangray_18_wilde_64kb.mp3'], 'length': 11528387},
#              {'path': ['Content', 'pictureofdoriangray_19_wilde_64kb.mp3'], 'length': 11443750},
#              {'path': ['Content', 'pictureofdoriangray_20_wilde_64kb.mp3'], 'length': 7585987},
#              {'path': ['Description.txt'], 'length': 1296},
#              {'path': ['License.txt'], 'length': 51}]

#     pudb.set_trace()
#     piece_length = 131072
#     want_files = [0,2]

#     for j in range(32):
#         # Want to find the rightmost file that begins at or before byte_index
#         write_file = get_next_want_file(byte_index=j*piece_length,
#                                         file_starts=file_starts,
#                                         files=files,
#                                         want_file_pos=want_files,
#                                         block='0'*piece_length)

#         print('piece {} starts at byte {} '
#               'and maps to {}'.format(j, j * piece_length, write_file))

# All the pieces for the first three files:
#
# piece 0 starts at byte 0 and maps to somefile(name='Content/174-h.htm')
# piece 1 starts at byte 131072 and maps to somefile(name='Content/174-h.htm')
# piece 2 starts at byte 262144 and maps to somefile(name='Content/174-h.htm')
# piece 3 starts at byte 393216 and maps to somefile(name='Content/174-h.htm')
# piece 4 starts at byte 524288 and maps to somefile(name='Content/174.txt')
# piece 5 starts at byte 655360 and maps to somefile(name='Content/174.txt')
# piece 6 starts at byte 786432 and maps to somefile(name='Content/174.txt')
# piece 7 starts at byte 917504 and maps to somefile(name='Content/174.txt')
# piece 8 starts at byte 1048576 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 9 starts at byte 1179648 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 10 starts at byte 1310720 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 11 starts at byte 1441792 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 12 starts at byte 1572864 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 13 starts at byte 1703936 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 14 starts at byte 1835008 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 15 starts at byte 1966080 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 16 starts at byte 2097152 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 17 starts at byte 2228224 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 18 starts at byte 2359296 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 19 starts at byte 2490368 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 20 starts at byte 2621440 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 21 starts at byte 2752512 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 22 starts at byte 2883584 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 23 starts at byte 3014656 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 24 starts at byte 3145728 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 25 starts at byte 3276800 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 26 starts at byte 3407872 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 27 starts at byte 3538944 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 28 starts at byte 3670016 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 29 starts at byte 3801088 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 30 starts at byte 3932160 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')
# piece 31 starts at byte 4063232 and maps to somefile(name='Content/pictureofdoriangray_00_wilde_64kb.mp3')

# if __name__ == '__main__':
#     test_get_write_file()
<|endoftext|>


<commit_before>import re
import functools
from collections import namedtuple
from ConfigParser import ConfigParser
from contextlib import contextmanager
from plumbum import ProcessExecutionError

__all__ = ['Repo']

class RepoConfig(object):
    '''
    This class provides an abstraction around repository configuration files
    '''

    def __init__(self, path):
        '''
        :param path:    Plumbum path to the config file
        '''
        self._config = ConfigParser()
        self._path = path
        if path.exists():
            self._config.readfp(path.open())
        else:
            self._config.add_section('paths')

    def AddRemote(self, name, destination):
        '''
        Adds a remote to the config, or overwrites if it already exists

        :param name:        The name of the remote
        :param destination: The destination path of the remote
        '''
        self._config.set('paths', name, destination)
        self._config.write(self._path.open('w'))

class Repo(object):
    '''
    This class provides an abstraction around running commands on a mercurial
    repository.  It can be used against either a local or remote repository
    depending on the ``machine`` parameter to the constructor.
    '''

    SummaryInfo = namedtuple('SummaryInfo', ['commit', 'mq'])
    CommitChangeInfo = namedtuple(
            'CommitChangeInfo',
            ['modified', 'unknown']
            )
    MqAppliedInfo = namedtuple('MqAppliedInfo', ['applied', 'unapplied'])

    # Template Parameter for hg log-style commands
    HgTemplateParam = '{node}\\t{desc|firstline}\\n'

    # Contains details of a changeset
    ChangesetInfo = namedtuple('ChangesetInfo', ['hash', 'desc'])
    ChangesetInfoRegexp = re.compile(r'^(?P<hash>\w+)\t(?P<desc>.*)$')

    def __init__(self, machine, remote=None):
        '''
        :param machine:     The plumbum machine object to use
                            (can be a local machine or remote machine)
        :param remote:      The name of the remote repo to be used by
                            push, pull and other operations.
        '''
        self.machine = machine
        self.hg = self.machine['hg']
        self.remote = remote
        self._currentRev = self._branch = None
        self.prevLevel = None

    @contextmanager
    def CleanMq(self):
        '''
        Returns a context manager that keeps the mq repository clean
        for it's lifetime
        '''
        revertTo = self.lastAppliedPatch
        self.PopPatch()
        yield
        if revertTo:
            self.PushPatch(revertTo)

    def _CleanMq(func):
        '''
        Decorator that ensures a function is always run with no patches applied
        Should only be applied on Repo methods

        :params func:   The function to decorate
        '''
        @functools.wraps(func)
        def InnerFunc(self, *pargs):
            with self.CleanMq():
                return func(self, *pargs)
        return InnerFunc

    @property
    def summary(self):
        '''
        Gets info from hg summary

        :return:    A :class:`SummaryInfo` containing :class:`CommitChangeInfo`
                    & :class:`MqAppliedInfo`
        '''
        commitData = Repo.CommitChangeInfo(0, 0)
        mqData = Repo.MqAppliedInfo(0, 0)
        commitRegexp = re.compile(
                r'^commit:\s+((\d+) modified(, (\d+) unknown)?)?'
                )
        mqRegexp = re.compile(
                r'^mq:\s+((\d+) applied,?\s*)?((\d+) unapplied)?'
                )
        lines = self.hg('summary').splitlines()
        for line in lines:
            match = commitRegexp.search( line )
            if match:
                commitData = Repo.CommitChangeInfo(
                        int(match.group(2) or 0), int(match.group(4) or 0)
                        )
            match = mqRegexp.search( line )
            if match:
                mqData = Repo.MqAppliedInfo(
                        int(match.group(2) or 0), int(match.group(4) or 0)
                        )
        return Repo.SummaryInfo(commitData, mqData)

    @property
    def currentRev(self):
        '''
        Gets the current revision
        This property is cached, so it may be out of date

        :returns:   A string containing the current revision hash
        '''
        if not self._currentRev:
            self._CheckCurrentRev()
        return self._currentRev

    @property
    def branch(self):
        '''
        Gets the current branch
        This property is cached, so it may be out of date

        :returns:   A string containing the current branch name
        '''
        if not self._branch:
            self._CheckCurrentRev()
        return self._branch

    @_CleanMq
    def _CheckCurrentRev( self ):
        ''' Gets the current revision and branch and stores it '''

        revMatch = re.search(
            r'^(\w{12})\+?\s+(.*)\s*$',
            self.hg('id', '-i', '-b')
            )
        if revMatch is None:
            raise Exception("Could not get current revision using hg id")
        self._currentRev, self._branch = revMatch.group(1, 2)

    def _RunListCommand(self, command, headerLines=0):
        '''
        Runs an hg command that gets a list (outgoing, incoming etc.)

        :param command:     The plumbum command object to run
        :param headerLines: The number of lines to chop off the top of
                            the output
        '''
        try:
            lines = command().splitlines()
            if headerLines == 0:
                return lines
            if len(lines) < headerLines:
                raise Exception("Unexpected number of lines from hg command")
            return lines[headerLines:]
        except ProcessExecutionError as e:
            if e.retcode != 1:
                # retcode of 1 just means there's nothing in the list,
                # so ignore it
                raise
        return []

    def _GetChangesetInfoList(self, *pargs, **kwargs):
        '''
        Utility function that calls _RunListCommand and filters the results
        through ChangesetInfoRegexp

        :returns:   A list of :class:`ChangesetInfo`
        '''
        lines = self._RunListCommand(*pargs, **kwargs)
        matches = (self.ChangesetInfoRegexp.match(line) for line in lines)
        return [self.ChangesetInfo(**match.groupdict()) for match in matches]

    @property
    @_CleanMq
    def outgoings(self):
        '''
        Gets the outgoing changesets to `self.remote`

        :returns:   A list containing :class:`ChangesetInfo` that represent
                    the current outgoing changesets
        '''
        assert self.remote
        return self._GetChangesetInfoList(
                self.hg[ 'outgoing', '-b', self.branch, '-r', self.currentRev,
                         '--template', self.HgTemplateParam, self.remote
                         ],
                headerLines=2
                )

    @property
    @_CleanMq
    def incomings(self):
        '''
        Gets the incoming changesets from `self.remote`

        :returns:   A list containing :class:`ChangesetInfo` that represent
                    the current incoming changesets
        '''
        assert self.remote
        return self._GetChangesetInfoList(
                self.hg[ 'incoming', '-b', self.branch,
                         '--template', self.HgTemplateParam, self.remote
                         ],
                headerLines=2
                )

    @property
    def lastAppliedPatch(self):
        '''
        Gets the last applied mq patch (if there is one)

        :returns: A single mq patch name (or None)
        '''
        try:
            return self.hg('qtop').strip()
        except ProcessExecutionError as e:
            if e.retcode not in [1, 255]:
                # 1 means no patches
                # 255 means mq is probably disabled
                raise
        return None

    @_CleanMq
    def PushToRemote(self):
        ''' Pushes to the remote repository at `self.remote`'''
        assert self.remote
        self.hg('push', '-b', self.branch, '-r', self.currentRev, self.remote)

    def PushMqToRemote(self):
        ''' Pushes the mq repo to the remote at `self.remote` '''
        assert self.remote
        try:
            self.hg('push', '--mq', self.remote)
        except ProcessExecutionError as e:
            if e.retcode != 1:
                #1 just means there's no outgoings
                raise

    def PopPatch(self, patch=None):
        '''
        Pops mq patch(es)

        :param patch:   Name of the patch to pop to.
                        If None, all will be popped
        '''
        # Check there are some patches applied
        level = self.lastAppliedPatch
        if level:
            if patch is None:
                patch = '-a'
            self.hg('qpop', patch)

    def PushPatch(self, patch=None):
        '''
        Pushes mq patch(es)

        :param patch:   Name of the patch to push to.
                        If None, all will be pushed
        '''
        if patch is None:
            patch = '-a'
        self.hg('qpush', patch)

    @_CleanMq
    def Strip(self, changesets):
        '''
        Strips changesets from this repository

        :param changesets:  A list of :class:`ChangesetInfo`
                            representing the changesets to strip
        '''
        self.hg('strip', *[cs.hash for cs in changesets])

    @_CleanMq
    def Update(self, changeset):
        '''
        Updates to a specific changeset

        :param changeset:   A changeset hash string, or
                            :class:`ChangesetInfo` representing
                            the changeset to update to
        '''
        if isinstance(changeset, self.ChangesetInfo):
            changeset = changeset.hash
        self.hg('update', changeset)

    def UpdateMq(self):
        '''
        Updates the mq repository to tip
        '''
        self.hg('update', '--mq')

    def RefreshMq(self):
        '''
        Refreshes the current mq patch
        '''
        self.hg('qrefresh')

    def CommitMq(self, msg=None):
        '''
        Commits the mq repository

        :param msg:     An optional commit message
        '''
        if not msg:
            msg = 'synchg-commit'
        try:
            self.hg('commit', '--mq', '-m', msg)
        except ProcessExecutionError as e:
            if e.retcode != 1:
                #1 just means there's no changes
                raise

    @_CleanMq
    def Clone(self, destination, createRemote=True):
        '''
        Clones the repository to a different location

        :param destination:     The destination clone path
        :param createRemote:    If set a remote will be created in the local
                                hgrc with the name the class was initialised
                                with.
        '''
        remoteName = self.remote if createRemote else None
        self._DoClone(destination, remoteName)
        patches = self.machine.cwd / '.hg' / 'patches'
        if patches.exists():
            # Clone mq repository
            # TODO: Possibly do a sanity check and call hg init --mq if needed
            with self.machine.cwd(patches):
                mqdest = destination + '/.hg/patches'
                self._DoClone(mqdest, remoteName)

    def _DoClone(self, destination, remoteName):
        '''
        Actually performs a clone operation

        :param destination:     The destination clone path
        :param remoteName:      The name of the remote to create (if any)
        '''
        self.hg('clone', '.', destination)
        if remoteName:
            configPath = self.machine.cwd / '.hg' / 'hgrc'
            config = RepoConfig(configPath)
            config.AddRemote(remoteName, destination)
<commit_msg>Added config property to Repo class<commit_after>import re
import functools
from collections import namedtuple
from ConfigParser import ConfigParser
from contextlib import contextmanager
from plumbum import ProcessExecutionError

__all__ = ['Repo']

class RepoConfig(object):
    '''
    This class provides an abstraction around repository configuration files
    '''

    def __init__(self, path):
        '''
        :param path:    Plumbum path to the repository
        '''
        self._config = ConfigParser()
        self._path = path / '.hg' / 'hgrc'
        if path.exists():
            self._config.readfp(path.open())
        else:
            self._config.add_section('paths')

    def AddRemote(self, name, destination):
        '''
        Adds a remote to the config, or overwrites if it already exists

        :param name:        The name of the remote
        :param destination: The destination path of the remote
        '''
        self._config.set('paths', name, destination)
        self._config.write(self._path.open('w'))

class Repo(object):
    '''
    This class provides an abstraction around running commands on a mercurial
    repository.  It can be used against either a local or remote repository
    depending on the ``machine`` parameter to the constructor.
    '''

    SummaryInfo = namedtuple('SummaryInfo', ['commit', 'mq'])
    CommitChangeInfo = namedtuple(
            'CommitChangeInfo',
            ['modified', 'unknown']
            )
    MqAppliedInfo = namedtuple('MqAppliedInfo', ['applied', 'unapplied'])

    # Template Parameter for hg log-style commands
    HgTemplateParam = '{node}\\t{desc|firstline}\\n'

    # Contains details of a changeset
    ChangesetInfo = namedtuple('ChangesetInfo', ['hash', 'desc'])
    ChangesetInfoRegexp = re.compile(r'^(?P<hash>\w+)\t(?P<desc>.*)$')

    def __init__(self, machine, remote=None):
        '''
        :param machine:     The plumbum machine object to use
                            (can be a local machine or remote machine)
        :param remote:      The name of the remote repo to be used by
                            push, pull and other operations.
        '''
        self.machine = machine
        self.hg = self.machine['hg']
        self.remote = remote
        self._currentRev = self._branch = None
        self.prevLevel = None
        self._config = self._mqconfig = None

    @contextmanager
    def CleanMq(self):
        '''
        Returns a context manager that keeps the mq repository clean
        for it's lifetime
        '''
        revertTo = self.lastAppliedPatch
        self.PopPatch()
        yield
        if revertTo:
            self.PushPatch(revertTo)

    def _CleanMq(func):
        '''
        Decorator that ensures a function is always run with no patches applied
        Should only be applied on Repo methods

        :params func:   The function to decorate
        '''
        @functools.wraps(func)
        def InnerFunc(self, *pargs):
            with self.CleanMq():
                return func(self, *pargs)
        return InnerFunc

    @property
    def summary(self):
        '''
        Gets info from hg summary

        :return:    A :class:`SummaryInfo` containing :class:`CommitChangeInfo`
                    & :class:`MqAppliedInfo`
        '''
        commitData = Repo.CommitChangeInfo(0, 0)
        mqData = Repo.MqAppliedInfo(0, 0)
        commitRegexp = re.compile(
                r'^commit:\s+((\d+) modified(, (\d+) unknown)?)?'
                )
        mqRegexp = re.compile(
                r'^mq:\s+((\d+) applied,?\s*)?((\d+) unapplied)?'
                )
        lines = self.hg('summary').splitlines()
        for line in lines:
            match = commitRegexp.search( line )
            if match:
                commitData = Repo.CommitChangeInfo(
                        int(match.group(2) or 0), int(match.group(4) or 0)
                        )
            match = mqRegexp.search( line )
            if match:
                mqData = Repo.MqAppliedInfo(
                        int(match.group(2) or 0), int(match.group(4) or 0)
                        )
        return Repo.SummaryInfo(commitData, mqData)

    @property
    def currentRev(self):
        '''
        Gets the current revision
        This property is cached, so it may be out of date

        :returns:   A string containing the current revision hash
        '''
        if not self._currentRev:
            self._CheckCurrentRev()
        return self._currentRev

    @property
    def branch(self):
        '''
        Gets the current branch
        This property is cached, so it may be out of date

        :returns:   A string containing the current branch name
        '''
        if not self._branch:
            self._CheckCurrentRev()
        return self._branch

    @property
    def config(self):
        '''
        Gets the configuration for this repository

        :returns: A ``RepoConfig`` class
        '''
        if not self._config:
            self._config = RepoConfig(self.machine.cwd)
        return self._config

    @property
    def mqconfig(self):
        '''
        Gets the configuration for the mq repository

        :returns A ``RepoConfig`` class
        '''
        if not self._mqconfig:
            self._mqconfig = RepoConfig(self.machine.cwd / '.hg' / 'patches')
        return self._mqconfig

    @_CleanMq
    def _CheckCurrentRev( self ):
        ''' Gets the current revision and branch and stores it '''

        revMatch = re.search(
            r'^(\w{12})\+?\s+(.*)\s*$',
            self.hg('id', '-i', '-b')
            )
        if revMatch is None:
            raise Exception("Could not get current revision using hg id")
        self._currentRev, self._branch = revMatch.group(1, 2)

    def _RunListCommand(self, command, headerLines=0):
        '''
        Runs an hg command that gets a list (outgoing, incoming etc.)

        :param command:     The plumbum command object to run
        :param headerLines: The number of lines to chop off the top of
                            the output
        '''
        try:
            lines = command().splitlines()
            if headerLines == 0:
                return lines
            if len(lines) < headerLines:
                raise Exception("Unexpected number of lines from hg command")
            return lines[headerLines:]
        except ProcessExecutionError as e:
            if e.retcode != 1:
                # retcode of 1 just means there's nothing in the list,
                # so ignore it
                raise
        return []

    def _GetChangesetInfoList(self, *pargs, **kwargs):
        '''
        Utility function that calls _RunListCommand and filters the results
        through ChangesetInfoRegexp

        :returns:   A list of :class:`ChangesetInfo`
        '''
        lines = self._RunListCommand(*pargs, **kwargs)
        matches = (self.ChangesetInfoRegexp.match(line) for line in lines)
        return [self.ChangesetInfo(**match.groupdict()) for match in matches]

    @property
    @_CleanMq
    def outgoings(self):
        '''
        Gets the outgoing changesets to `self.remote`

        :returns:   A list containing :class:`ChangesetInfo` that represent
                    the current outgoing changesets
        '''
        assert self.remote
        return self._GetChangesetInfoList(
                self.hg[ 'outgoing', '-b', self.branch, '-r', self.currentRev,
                         '--template', self.HgTemplateParam, self.remote
                         ],
                headerLines=2
                )

    @property
    @_CleanMq
    def incomings(self):
        '''
        Gets the incoming changesets from `self.remote`

        :returns:   A list containing :class:`ChangesetInfo` that represent
                    the current incoming changesets
        '''
        assert self.remote
        return self._GetChangesetInfoList(
                self.hg[ 'incoming', '-b', self.branch,
                         '--template', self.HgTemplateParam, self.remote
                         ],
                headerLines=2
                )

    @property
    def lastAppliedPatch(self):
        '''
        Gets the last applied mq patch (if there is one)

        :returns: A single mq patch name (or None)
        '''
        try:
            return self.hg('qtop').strip()
        except ProcessExecutionError as e:
            if e.retcode not in [1, 255]:
                # 1 means no patches
                # 255 means mq is probably disabled
                raise
        return None

    @_CleanMq
    def PushToRemote(self):
        ''' Pushes to the remote repository at `self.remote`'''
        assert self.remote
        self.hg('push', '-b', self.branch, '-r', self.currentRev, self.remote)

    def PushMqToRemote(self):
        ''' Pushes the mq repo to the remote at `self.remote` '''
        assert self.remote
        try:
            self.hg('push', '--mq', self.remote)
        except ProcessExecutionError as e:
            if e.retcode != 1:
                #1 just means there's no outgoings
                raise

    def PopPatch(self, patch=None):
        '''
        Pops mq patch(es)

        :param patch:   Name of the patch to pop to.
                        If None, all will be popped
        '''
        # Check there are some patches applied
        level = self.lastAppliedPatch
        if level:
            if patch is None:
                patch = '-a'
            self.hg('qpop', patch)

    def PushPatch(self, patch=None):
        '''
        Pushes mq patch(es)

        :param patch:   Name of the patch to push to.
                        If None, all will be pushed
        '''
        if patch is None:
            patch = '-a'
        self.hg('qpush', patch)

    @_CleanMq
    def Strip(self, changesets):
        '''
        Strips changesets from this repository

        :param changesets:  A list of :class:`ChangesetInfo`
                            representing the changesets to strip
        '''
        self.hg('strip', *[cs.hash for cs in changesets])

    @_CleanMq
    def Update(self, changeset):
        '''
        Updates to a specific changeset

        :param changeset:   A changeset hash string, or
                            :class:`ChangesetInfo` representing
                            the changeset to update to
        '''
        if isinstance(changeset, self.ChangesetInfo):
            changeset = changeset.hash
        self.hg('update', changeset)

    def UpdateMq(self):
        '''
        Updates the mq repository to tip
        '''
        self.hg('update', '--mq')

    def RefreshMq(self):
        '''
        Refreshes the current mq patch
        '''
        self.hg('qrefresh')

    def CommitMq(self, msg=None):
        '''
        Commits the mq repository

        :param msg:     An optional commit message
        '''
        if not msg:
            msg = 'synchg-commit'
        try:
            self.hg('commit', '--mq', '-m', msg)
        except ProcessExecutionError as e:
            if e.retcode != 1:
                #1 just means there's no changes
                raise

    @_CleanMq
    def Clone(self, destination, createRemote=True):
        '''
        Clones the repository to a different location

        :param destination:     The destination clone path
        :param createRemote:    If set a remote will be created in the local
                                hgrc with the name the class was initialised
                                with.
        '''
        remoteName = self.remote if createRemote else None
        self._DoClone(destination, remoteName)
        patches = self.machine.cwd / '.hg' / 'patches'
        if patches.exists():
            # Clone mq repository
            # TODO: Possibly do a sanity check and call hg init --mq if needed
            with self.machine.cwd(patches):
                mqdest = destination + '/.hg/patches'
                self._DoClone(mqdest, remoteName)

    def _DoClone(self, destination, remoteName):
        '''
        Actually performs a clone operation

        :param destination:     The destination clone path
        :param remoteName:      The name of the remote to create (if any)
        '''
        self.hg('clone', '.', destination)
        if remoteName:
            configPath = self.machine.cwd / '.hg' / 'hgrc'
            config = RepoConfig(configPath)
            config.AddRemote(remoteName, destination)
<|endoftext|>


<commit_before>import md5
import random
import simplejson
import logging
import time

from django.conf import settings

from google.appengine.ext import db

__all__ = ['User', 'Group', 'Status', 'FullStatusError']

def serialize(obj):
    start = time.time()
    out = simplejson.dumps(obj)
    logging.info('serializing time: %f', time.time() - start)
    return out

def deserialize(obj):
    start = time.time()
    out = simplejson.loads(obj)
    logging.info('deserializing time: %f', time.time() - start)
    return out

class User(db.Model):
    username = db.StringProperty()
    password_hash = db.StringProperty()

    # pretend to be a property
    realm = 'Synk'

    def set_password(self, password):
        # HTTP digest user-realm-pass hash
        a1 = '%s:%s:%s' % (self.username, self.realm, password)
        ha1 = md5.new(a1)
        self.password_hash = ha1.hexdigest()

    @staticmethod
    def by_username(username):
        try:
            user = db.GqlQuery('select * from User where username = :1 limit 1', username)[0]
            return user
        except IndexError:
            return None

class Group(db.Model):
    user = db.ReferenceProperty(User)

    # the first two characters of the key of
    # each item in the status_map (see Status)
    prefix = db.StringProperty()

    # total per-user capacity is 16 ** PREFIX_LEN * 11000 * 1000
    # ... so about 176M (with theoretical full capacity and no
    # key collisions... not sure if this is realistic)
    #
    # for users with smaller collections (probably most users)
    # setting this too high causes lots of datastore overhead
    # because we're forced to do a lot of looping in the
    # /status handler. leave it at 1 for now unless it
    # becomes apparent that this is limiting
    PREFIX_LEN = 1

    @staticmethod
    def for_user(user):
        start = time.time()
        groups = db.GqlQuery('select * from Group where user = :1', user)
        end = time.time()
        logging.info("Group.for_user: %f", end - start)
        return groups

    @staticmethod
    def for_user_prefixes(user, prefixes):
        start = time.time()
        groups = db.GqlQuery('select * from Group where user = :1', user)
        groups = [group for group in groups if group.prefix in set(prefixes)]
        end = time.time()
        logging.info("Group.for_user_prefixes: %f", end - start)
        return groups

    def get_statuses(self):
        start = time.time()
        statuses = db.GqlQuery('select * from Status where group = :1 order by fill_factor desc', self)
        end = time.time()
        logging.info("Group.get_statuses: %f", end - start)
        return [s for s in statuses]

class Status(db.Model):
    STATUS_UNREAD = 0
    STATUS_READ = 1

    STATUS_MAP = {}
    STATUS_MAP['unread'] = STATUS_UNREAD
    STATUS_MAP['read'] = STATUS_READ

    INVERSE_STATUS_MAP = dict([(value, key) for key, value in STATUS_MAP.iteritems()])
    STATUS_VALUES = set(STATUS_MAP.values())

    group = db.ReferenceProperty(Group)

    # serialized python dictionary mapping item id
    # to status key
    status_map_serialized = db.TextProperty()
    status_map = None

    # max_size is based on max_len of 900000 bytes
    # with a constant per-entry size of 79 bytes
    # (this is true when serializing as JSON)
    max_size = 11000
    max_len = 900000
    fill_factor = db.FloatProperty()

    def init_status_map(self):
        if self.status_map is None and self.status_map_serialized:
            self.status_map = deserialize(self.status_map_serialized)
            logging.info("%d items, %d bytes, %f fill factor", len(self.status_map), len(self.status_map_serialized), self.fill_factor)
        elif self.status_map is None:
            self.status_map = {}

    def has_item(self, id):
        self.init_status_map()
        return id in self.status_map

    def get_item(self, id):
        self.init_status_map()
        try:
            return self.status_map[id]
        except KeyError:
            return None

    def is_full(self):
        return settings.PROFILING and len(serialize(self.status_map)) >= self.max_len \
            or len(self.status_map) >= self.max_size

    def set_item(self, id, value):
        self.init_status_map()
        if not self.has_item(id) and self.is_full():
            logging.warn("raising FullStatusError with %d items", len(self.status_map))
            raise FullStatusError('Status is full')
        self.status_map[id] = value

    def del_item(self, id):
        self.init_status_map()
        del self.status_map[id]

    def put(self):
        self.status_map_serialized = serialize(self.status_map)
        self.fill_factor = float(len(self.status_map_serialized)) / float(self.max_len)
        db.Model.put(self)
        if settings.PROFILING:
            logging.info("%d items, %d bytes, %f fill factor", len(self.status_map), len(self.status_map_serialized), self.fill_factor)

class FullStatusError(Exception):
    """Raised when the status is at its max_size already"""
    pass
<commit_msg>use more compact json serialization<commit_after>import md5
import random
import simplejson
import logging
import time

from django.conf import settings

from google.appengine.ext import db

__all__ = ['User', 'Group', 'Status', 'FullStatusError']

def serialize(obj):
    start = time.time()
    out = simplejson.dumps(obj, separators=[',', ':'])
    logging.info('serializing time: %f', time.time() - start)
    return out

def deserialize(obj):
    start = time.time()
    out = simplejson.loads(obj)
    logging.info('deserializing time: %f', time.time() - start)
    return out

class User(db.Model):
    username = db.StringProperty()
    password_hash = db.StringProperty()

    # pretend to be a property
    realm = 'Synk'

    def set_password(self, password):
        # HTTP digest user-realm-pass hash
        a1 = '%s:%s:%s' % (self.username, self.realm, password)
        ha1 = md5.new(a1)
        self.password_hash = ha1.hexdigest()

    @staticmethod
    def by_username(username):
        try:
            user = db.GqlQuery('select * from User where username = :1 limit 1', username)[0]
            return user
        except IndexError:
            return None

class Group(db.Model):
    user = db.ReferenceProperty(User)

    # the first two characters of the key of
    # each item in the status_map (see Status)
    prefix = db.StringProperty()

    # total per-user capacity is 16 ** PREFIX_LEN * 11000 * 1000
    # ... so about 176M (with theoretical full capacity and no
    # key collisions... not sure if this is realistic)
    #
    # for users with smaller collections (probably most users)
    # setting this too high causes lots of datastore overhead
    # because we're forced to do a lot of looping in the
    # /status handler. leave it at 1 for now unless it
    # becomes apparent that this is limiting
    PREFIX_LEN = 1

    @staticmethod
    def for_user(user):
        start = time.time()
        groups = db.GqlQuery('select * from Group where user = :1', user)
        end = time.time()
        logging.info("Group.for_user: %f", end - start)
        return groups

    @staticmethod
    def for_user_prefixes(user, prefixes):
        start = time.time()
        groups = db.GqlQuery('select * from Group where user = :1', user)
        groups = [group for group in groups if group.prefix in set(prefixes)]
        end = time.time()
        logging.info("Group.for_user_prefixes: %f", end - start)
        return groups

    def get_statuses(self):
        start = time.time()
        statuses = db.GqlQuery('select * from Status where group = :1 order by fill_factor desc', self)
        end = time.time()
        logging.info("Group.get_statuses: %f", end - start)
        return [s for s in statuses]

class Status(db.Model):
    STATUS_UNREAD = 0
    STATUS_READ = 1

    STATUS_MAP = {}
    STATUS_MAP['unread'] = STATUS_UNREAD
    STATUS_MAP['read'] = STATUS_READ

    INVERSE_STATUS_MAP = dict([(value, key) for key, value in STATUS_MAP.iteritems()])
    STATUS_VALUES = set(STATUS_MAP.values())

    group = db.ReferenceProperty(Group)

    # serialized python dictionary mapping item id
    # to status key
    status_map_serialized = db.TextProperty()
    status_map = None

    # max_size is based on max_len of 900000 bytes
    # with a constant per-entry size of 79 bytes
    # (this is true when serializing as JSON)
    max_size = 11000
    max_len = 900000
    fill_factor = db.FloatProperty()

    def init_status_map(self):
        if self.status_map is None and self.status_map_serialized:
            self.status_map = deserialize(self.status_map_serialized)
            logging.info("%d items, %d bytes, %f fill factor", len(self.status_map), len(self.status_map_serialized), self.fill_factor)
        elif self.status_map is None:
            self.status_map = {}

    def has_item(self, id):
        self.init_status_map()
        return id in self.status_map

    def get_item(self, id):
        self.init_status_map()
        try:
            return self.status_map[id]
        except KeyError:
            return None

    def is_full(self):
        return settings.PROFILING and len(serialize(self.status_map)) >= self.max_len \
            or len(self.status_map) >= self.max_size

    def set_item(self, id, value):
        self.init_status_map()
        if not self.has_item(id) and self.is_full():
            logging.warn("raising FullStatusError with %d items", len(self.status_map))
            raise FullStatusError('Status is full')
        self.status_map[id] = value

    def del_item(self, id):
        self.init_status_map()
        del self.status_map[id]

    def put(self):
        self.status_map_serialized = serialize(self.status_map)
        self.fill_factor = float(len(self.status_map_serialized)) / float(self.max_len)
        db.Model.put(self)
        if settings.PROFILING:
            logging.info("%d items, %d bytes, %f fill factor", len(self.status_map), len(self.status_map_serialized), self.fill_factor)

class FullStatusError(Exception):
    """Raised when the status is at its max_size already"""
    pass
<|endoftext|>