Added MemoryDB including initial test, moved ZippedShaWriter into stream module, it was just a test helper previously

Author: Byron

base.py

@@ -41,6 +41,16 @@ def __init__(self, *args):
 	def sha(self):
 		return self[0]
 
+	@property
+	def hexsha(self):
+		""":return: our sha, hex encoded, 40 bytes"""
+		return to_hex_sha(self[0])
+		
+	@property
+	def binsha(self):
+		""":return: our sha as binary, 20 bytes"""
+		return to_bin_sha(self[0])
+		
 	@property
 	def type(self):
 		return self[1]

db/__init__.py

@@ -1,6 +1,7 @@
 
 from base import *
 from loose import *
+from mem import *
 from pack import *
 from git import *
 from ref import *

db/mem.py

@@ -0,0 +1,80 @@
+"""Contains the MemoryDatabase implementation"""
+from loose import LooseObjectDB
+from base import (
+						ObjectDBR, 
+						ObjectDBW
+					)
+
+from gitdb.base import OStream
+from gitdb.util import to_bin_sha
+from gitdb.exc import (
+						BadObject,
+						UnsupportedOperation
+						)
+from gitdb.stream import (
+							ZippedStoreShaWriter,
+							DecompressMemMapReader,
+						)
+
+__all__ = ("MemoryDB", )
+
+class MemoryDB(ObjectDBR, ObjectDBW):
+	"""A memory database stores everything to memory, providing fast IO and object
+	retrieval. It should be used to buffer results and obtain SHAs before writing
+	it to the actual physical storage, as it allows to query whether object already
+	exists in the target storage before introducing actual IO
+	
+	:note: memory is currently not threadsafe, hence the async methods cannot be used
+		for storing"""
+	
+	def __init__(self):
+		super(MemoryDB, self).__init__()
+		self._db = LooseObjectDB("path/doesnt/matter")
+		
+		# maps 20 byte shas to their OStream objects
+		self._cache = dict()
+		
+	def set_ostream(self, stream):
+		raise UnsupportedOperation("MemoryDB's always stream into memory")
+		
+	def store(self, istream):
+		zstream = ZippedStoreShaWriter()
+		self._db.set_ostream(zstream)
+		
+		istream = self._db.store(istream)
+		zstream.close()		# close to flush
+		zstream.seek(0)
+		
+		# don't provide a size, the stream is written in object format, hence the 
+		# header needs decompression
+		decomp_stream = DecompressMemMapReader(zstream.getvalue(), close_on_deletion=False) 
+		self._cache[istream.binsha] = OStream(istream.sha, istream.type, istream.size, decomp_stream)
+		
+		return istream
+		
+	def store_async(self, reader):
+		raise UnsupportedOperation("MemoryDBs cannot currently be used for async write access")
+	
+	def has_object(self, sha):
+		return to_bin_sha(sha) in self._cache
+
+	def info(self, sha):
+		# we always return streams, which are infos as well
+		return self.stream(sha)
+	
+	def stream(self, sha):
+		sha = to_bin_sha(sha)
+		try:
+			ostream = self._cache[sha]
+			# rewind stream for the next one to read
+			ostream.stream.seek(0)
+			return ostream
+		except KeyError:
+			raise BadObject(sha)
+		# END exception handling
+	
+	def size(self):
+		return len(self._cache)
+		
+	def sha_iter(self):
+		return self._cache.iterkeys()

stream.py

@@ -499,6 +499,39 @@ def sha(self, as_hex = False):
 
 	#} END interface 
 
+
+class ZippedStoreShaWriter(Sha1Writer):
+	"""Remembers everything someone writes to it and generates a sha"""
+	__slots__ = ('buf', 'zip')
+	def __init__(self):
+		Sha1Writer.__init__(self)
+		self.buf = StringIO()
+		self.zip = zlib.compressobj(zlib.Z_BEST_SPEED)
+	
+	def __getattr__(self, attr):
+		return getattr(self.buf, attr)
+	
+	def write(self, data):
+		alen = Sha1Writer.write(self, data)
+		self.buf.write(self.zip.compress(data))
+		return alen
+		
+	def close(self):
+		self.buf.write(self.zip.flush())
+		
+	def seek(self, offset, whence=os.SEEK_SET):
+		"""Seeking currently only supports to rewind written data
+		Multiple writes are not supported"""
+		if offset != 0 or whence != os.SEEK_SET:
+			raise ValueError("Can only seek to position 0")
+		# END handle offset
+		self.buf.seek(0)
+		
+	def getvalue(self):
+		""":return: string value from the current stream position to the end"""
+		return self.buf.getvalue()
+
+
 class FDCompressedSha1Writer(Sha1Writer):
 	"""Digests data written to it, making the sha available, then compress the 
 	data and write it to the file descriptor

test/db/lib.py

@@ -20,6 +20,7 @@
 
 from async import IteratorReader
 from cStringIO import StringIO
+from struct import pack
 
 
 __all__ = ('TestDBBase', 'with_rw_directory', 'with_packs_rw', 'fixture_path')
@@ -29,9 +30,35 @@ class TestDBBase(TestBase):
 
 	# data
 	two_lines = "1234\nhello world"
-	
 	all_data = (two_lines, )
 
+	def _assert_object_writing_simple(self, db):
+		# write a bunch of objects and query their streams and info
+		null_objs = db.size()
+		ni = 250
+		for i in xrange(ni):
+			data = pack(">L", i)
+			istream = IStream(str_blob_type, len(data), StringIO(data))
+			new_istream = db.store(istream)
+			assert new_istream is istream
+			assert db.has_object(istream.sha)
+			
+			info = db.info(istream.sha)
+			assert isinstance(info, OInfo)
+			assert info.type == istream.type and info.size == istream.size
+			
+			stream = db.stream(istream.sha)
+			assert isinstance(stream, OStream)
+			assert stream.sha == info.sha and stream.type == info.type
+			assert stream.read() == data
+		# END for each item
+		
+		assert db.size() == null_objs + ni
+		shas = list(db.sha_iter())
+		assert len(shas) == db.size()
+		assert len(shas[0]) == 20
+		
+	
 	def _assert_object_writing(self, db):
 		"""General tests to verify object writing, compatible to ObjectDBW
 		:note: requires write access to the database"""

test/db/test_mem.py

@@ -0,0 +1,10 @@
+from lib import *
+from gitdb.db import MemoryDB
+		
+class TestMemoryDB(TestDBBase):
+	
+	def test_writing(self):
+		mdb = MemoryDB()
+		
+		# write data
+		self._assert_object_writing_simple(mdb)

test/lib.py

@@ -2,7 +2,11 @@
 from gitdb import (
 	OStream, 
 	)
-from gitdb.stream import Sha1Writer
+from gitdb.stream import ( 
+							Sha1Writer, 
+							ZippedStoreShaWriter
+						)
+
 from gitdb.util import zlib
 
 import sys
@@ -140,26 +144,5 @@ def _assert(self):
 		assert self.args
 		assert self.myarg
 
-
-class ZippedStoreShaWriter(Sha1Writer):
-	"""Remembers everything someone writes to it"""
-	__slots__ = ('buf', 'zip')
-	def __init__(self):
-		Sha1Writer.__init__(self)
-		self.buf = StringIO()
-		self.zip = zlib.compressobj(1)	# fastest
-	
-	def __getattr__(self, attr):
-		return getattr(self.buf, attr)
-	
-	def write(self, data):
-		alen = Sha1Writer.write(self, data)
-		self.buf.write(self.zip.compress(data))
-		return alen
-		
-	def close(self):
-		self.buf.write(self.zip.flush())
-
-
 #} END stream utilitiess
 

test/test_base.py

@@ -31,12 +31,12 @@ def test_streams(self):
 
 		# test pack info
 		# provides type_id
-		pinfo = OPackInfo(0, 1, blob_id, s)
+		pinfo = OPackInfo(0, blob_id, s)
 		assert pinfo.type == str_blob_type
 		assert pinfo.type_id == blob_id
 		assert pinfo.pack_offset == 0
 
-		dpinfo = ODeltaPackInfo(0, 1, blob_id, s, sha)
+		dpinfo = ODeltaPackInfo(0, blob_id, s, sha)
 		assert dpinfo.type == str_blob_type
 		assert dpinfo.type_id == blob_id
 		assert dpinfo.delta_info == sha