e7ff64716835249467a73f2a996789e3a91e24e0
[jelmer/dulwich-libgit2.git] / dulwich / object_store.py
1 # object_store.py -- Object store for git objects 
2 # Copyright (C) 2008 Jelmer Vernooij <jelmer@samba.org>
3
4 # This program is free software; you can redistribute it and/or
5 # modify it under the terms of the GNU General Public License
6 # as published by the Free Software Foundation; either version 2
7 # or (at your option) a later version of the License.
8
9 # This program is distributed in the hope that it will be useful,
10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12 # GNU General Public License for more details.
13
14 # You should have received a copy of the GNU General Public License
15 # along with this program; if not, write to the Free Software
16 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
17 # MA  02110-1301, USA.
18
19 import os
20 import tempfile
21 import urllib2
22
23 from dulwich.objects import (
24     hex_to_sha,
25     sha_to_hex,
26     ShaFile,
27     )
28 from dulwich.pack import (
29     iter_sha1, 
30     load_packs, 
31     write_pack,
32     write_pack_data,
33     write_pack_index_v2,
34     PackData, 
35     )
36
37 PACKDIR = 'pack'
38
39 class ObjectStore(object):
40
41     def __init__(self, path):
42         self.path = path
43         self._packs = None
44
45     def iter_shas(self, shas):
46         return ObjectStoreIterator(self, shas)
47
48     def pack_dir(self):
49         return os.path.join(self.path, PACKDIR)
50
51     def __contains__(self, sha):
52         # TODO: This can be more efficient
53         try:
54             self[sha]
55             return True
56         except KeyError:
57             return False
58
59     @property
60     def packs(self):
61         """List with pack objects."""
62         if self._packs is None:
63             self._packs = list(load_packs(self.pack_dir()))
64         return self._packs
65
66     def _get_shafile_path(self, sha):
67         dir = sha[:2]
68         file = sha[2:]
69         # Check from object dir
70         return os.path.join(self.path, dir, file)
71
72     def _get_shafile(self, sha):
73         path = self._get_shafile_path(sha)
74         if os.path.exists(path):
75           return ShaFile.from_file(path)
76         return None
77
78     def _add_shafile(self, sha, o):
79         path = self._get_shafile_path(sha)
80         f = os.path.open(path, 'w')
81         try:
82             f.write(o._header())
83             f.write(o._text)
84         finally:
85             f.close()
86
87     def get_raw(self, sha):
88         """Obtain the raw text for an object.
89         
90         :param sha: Sha for the object.
91         :return: tuple with object type and object contents.
92         """
93         for pack in self.packs:
94             if sha in pack:
95                 return pack.get_raw(sha, self.get_raw)
96         # FIXME: Are pack deltas ever against on-disk shafiles ?
97         ret = self._get_shafile(sha)
98         if ret is not None:
99             return ret.as_raw_string()
100         raise KeyError(sha)
101
102     def __getitem__(self, sha):
103         assert len(sha) == 40, "Incorrect length sha: %s" % str(sha)
104         ret = self._get_shafile(sha)
105         if ret is not None:
106             return ret
107         # Check from packs
108         type, uncomp = self.get_raw(sha)
109         return ShaFile.from_raw_string(type, uncomp)
110
111     def move_in_thin_pack(self, path):
112         """Move a specific file containing a pack into the pack directory.
113
114         :note: The file should be on the same file system as the 
115             packs directory.
116
117         :param path: Path to the pack file.
118         """
119         p = PackData(path)
120         temppath = os.path.join(self.pack_dir(), sha_to_hex(urllib2.randombytes(20))+".temppack")
121         write_pack(temppath, p.iterobjects(self.get_raw), len(p))
122         pack_sha = PackIndex(temppath+".idx").objects_sha1()
123         os.rename(temppath+".pack", 
124             os.path.join(self.pack_dir(), "pack-%s.pack" % pack_sha))
125         os.rename(temppath+".idx", 
126             os.path.join(self.pack_dir(), "pack-%s.idx" % pack_sha))
127
128     def move_in_pack(self, path):
129         """Move a specific file containing a pack into the pack directory.
130
131         :note: The file should be on the same file system as the 
132             packs directory.
133
134         :param path: Path to the pack file.
135         """
136         p = PackData(path)
137         entries = p.sorted_entries()
138         basename = os.path.join(self.pack_dir(), 
139             "pack-%s" % iter_sha1(entry[0] for entry in entries))
140         write_pack_index_v2(basename+".idx", entries, p.calculate_checksum())
141         os.rename(path, basename + ".pack")
142
143     def add_thin_pack(self):
144         """Add a new thin pack to this object store.
145
146         Thin packs are packs that contain deltas with parents that exist 
147         in a different pack.
148         """
149         fd, path = tempfile.mkstemp(dir=self.pack_dir(), suffix=".pack")
150         f = os.fdopen(fd, 'w')
151         def commit():
152             if os.path.getsize(path) > 0:
153                 self.move_in_thin_pack(path)
154         return f, commit
155
156     def add_pack(self):
157         """Add a new pack to this object store. 
158
159         :return: Fileobject to write to and a commit function to 
160             call when the pack is finished.
161         """
162         fd, path = tempfile.mkstemp(dir=self.pack_dir(), suffix=".pack")
163         f = os.fdopen(fd, 'w')
164         def commit():
165             if os.path.getsize(path) > 0:
166                 self.move_in_pack(path)
167         return f, commit
168
169     def add_objects(self, objects):
170         if len(objects) == 0:
171             return
172         f, commit = self.add_pack()
173         write_pack_data(f, objects, len(objects))
174         commit()
175
176
177 class ObjectImporter(object):
178
179     def __init__(self, count):
180         self.count = count
181
182     def add_object(self, object):
183         raise NotImplementedError(self.add_object)
184
185     def finish(self, object):
186         raise NotImplementedError(self.finish)
187
188
189 class ObjectIterator(object):
190
191     def iterobjects(self):
192         raise NotImplementedError(self.iterobjects)
193
194
195 class ObjectStoreIterator(ObjectIterator):
196
197     def __init__(self, store, shas):
198         self.store = store
199         self.shas = shas
200
201     def __iter__(self):
202         return ((self.store[sha], path) for sha, path in self.shas)
203
204     def iterobjects(self):
205         for o, path in self:
206             yield o
207
208     def __contains__(self, needle):
209         """Check if an object is present.
210
211         :param needle: SHA1 of the object to check for
212         """
213         # FIXME: This could be more efficient
214         for sha, path in self.shas:
215             if sha == needle:
216                 return True
217         return False
218
219     def __getitem__(self, key):
220         """Find an object by SHA1."""
221         return self.store[key]
222
223     def __len__(self):
224         """Return the number of objects."""
225         return len(self.shas)
226
227