diff --git a/plugins/__init__.py b/plugins/__init__.py index 382b74180..54e41981c 100644 --- a/plugins/__init__.py +++ b/plugins/__init__.py @@ -38,6 +38,7 @@ import os import re import csv import sys +import math import sets import time import random @@ -357,6 +358,135 @@ class PeriodicFileDownloader(object): world.threadsSpawned += 1 +class FlatfileDB(object): + def __init__(self, filename, maxSize=10**6): + self.filename = filename + try: + fd = file(self.filename) + strId = fd.readline().rstrip() + self.maxSize = len(strId) + self.currentId = int(strId) + except EnvironmentError, e: + # File couldn't be opened. + self.maxSize = int(math.log10(maxSize)) + self.currentId = 0 + self._incrementCurrentId() + + def serialize(self, record): + raise NotImplementedError + + def unserialize(self, s): + raise NotImplementedError + + def _canonicalId(self, id): + if id is not None: + return str(id).zfill(self.maxSize) + else: + return '-'*self.maxSize + + def _incrementCurrentId(self, fd=None): + fdWasNone = fd is None + if fdWasNone: + fd = file(self.filename, 'a') + fd.seek(0) + self.currentId += 1 + fd.write(self._canonicalId(self.currentId)) + fd.write('\n') + if fdWasNone: + fd.close() + + def _splitLine(self, line): + line = line.rstrip('\r\n') + (strId, strRecord) = line.split(':', 1) + return (strId, strRecord) + + def _joinLine(self, id, record): + return '%s:%s\n' % (self._canonicalId(id), self.serialize(record)) + + def addRecord(self, record): + line = self._joinLine(self.currentId, record) + try: + fd = file(self.filename, 'r+') + fd.seek(0, 2) # End. + fd.write(line) + self._incrementCurrentId(fd) + finally: + fd.close() + + def getRecord(self, id): + strId = self._canonicalId(id) + try: + fd = file(self.filename) + fd.readline() # First line, nextId. + for line in fd: + (lineId, strRecord) = self._splitLine(line) + if lineId == strId: + return self.unserialize(strRecord) + raise KeyError, id + finally: + fd.close() + + def setRecord(self, id, record): + strLine = self._joinLine(id, record) + try: + fd = file(self.filename, 'r+') + self.delRecord(id, fd) + fd.seek(0, 2) # End. + fd.write(strLine) + finally: + fd.close() + + def delRecord(self, id, fd=None): + fdWasNone = fd is None + strId = self._canonicalId(id) + try: + if fdWasNone: + fd = file(self.filename, 'r+') + fd.seek(0) + fd.readline() # First line, nextId + pos = fd.tell() + line = fd.readline() + while line: + (lineId, strRecord) = self._splitLine(line) + if lineId == strId: + fd.seek(pos) + fd.write(self._canonicalId(None)) + fd.seek(pos) + fd.readline() # Same line we just rewrote the id for. + pos = fd.tell() + line = fd.readline() + # We should be at the end. + finally: + if fdWasNone: + fd.close() + + def records(self): + fd = file(self.filename) + fd.readline() # First line, nextId. + for line in fd: + (strId, strRecord) = self._splitLine(line) + if not strId.startswith('-'): + yield (int(strId), self.unserialize(strRecord)) + fd.close() + + def vacuum(self): + infd = file(self.filename) + outfd = utils.transactionalFile(self.filename) + outfd.write(infd.readline()) # First line, nextId. + for line in infd: + if not line.startswith('-'): + outfd.write(line) + infd.close() + outfd.close() + + def flush(self): + pass # No-op, we maintain no open files. + + def close(self): + self.vacuum() # Should we do this? It should be fine. + + + _randomnickRe = re.compile(r'\$rand(?:om)?nick', re.I) _randomdateRe = re.compile(r'\$rand(?:om)?date', re.I) _randomintRe = re.compile(r'\$rand(?:omint)?', re.I)