##// END OF EJS Templates
infinitepush: replace `NamedTemporaryFile` with `pycompat.namedtempfile`...
Connor Sheehan -
r45752:a52bf967 stable
parent child Browse files
Show More
@@ -1,186 +1,183 b''
1 # This software may be used and distributed according to the terms of the
1 # This software may be used and distributed according to the terms of the
2 # GNU General Public License version 2 or any later version.
2 # GNU General Public License version 2 or any later version.
3
3
4 # based on bundleheads extension by Gregory Szorc <gps@mozilla.com>
4 # based on bundleheads extension by Gregory Szorc <gps@mozilla.com>
5
5
6 from __future__ import absolute_import
6 from __future__ import absolute_import
7
7
8 import abc
8 import abc
9 import os
9 import os
10 import subprocess
10 import subprocess
11 import tempfile
11 import tempfile
12
12
13 from mercurial.pycompat import open
13 from mercurial.pycompat import open
14 from mercurial import (
14 from mercurial import (
15 node,
15 node,
16 pycompat,
16 pycompat,
17 )
17 )
18 from mercurial.utils import (
18 from mercurial.utils import (
19 hashutil,
19 hashutil,
20 procutil,
20 procutil,
21 )
21 )
22
22
23 NamedTemporaryFile = tempfile.NamedTemporaryFile
24
25
23
26 class BundleWriteException(Exception):
24 class BundleWriteException(Exception):
27 pass
25 pass
28
26
29
27
30 class BundleReadException(Exception):
28 class BundleReadException(Exception):
31 pass
29 pass
32
30
33
31
34 class abstractbundlestore(object): # pytype: disable=ignored-metaclass
32 class abstractbundlestore(object): # pytype: disable=ignored-metaclass
35 """Defines the interface for bundle stores.
33 """Defines the interface for bundle stores.
36
34
37 A bundle store is an entity that stores raw bundle data. It is a simple
35 A bundle store is an entity that stores raw bundle data. It is a simple
38 key-value store. However, the keys are chosen by the store. The keys can
36 key-value store. However, the keys are chosen by the store. The keys can
39 be any Python object understood by the corresponding bundle index (see
37 be any Python object understood by the corresponding bundle index (see
40 ``abstractbundleindex`` below).
38 ``abstractbundleindex`` below).
41 """
39 """
42
40
43 __metaclass__ = abc.ABCMeta
41 __metaclass__ = abc.ABCMeta
44
42
45 @abc.abstractmethod
43 @abc.abstractmethod
46 def write(self, data):
44 def write(self, data):
47 """Write bundle data to the store.
45 """Write bundle data to the store.
48
46
49 This function receives the raw data to be written as a str.
47 This function receives the raw data to be written as a str.
50 Throws BundleWriteException
48 Throws BundleWriteException
51 The key of the written data MUST be returned.
49 The key of the written data MUST be returned.
52 """
50 """
53
51
54 @abc.abstractmethod
52 @abc.abstractmethod
55 def read(self, key):
53 def read(self, key):
56 """Obtain bundle data for a key.
54 """Obtain bundle data for a key.
57
55
58 Returns None if the bundle isn't known.
56 Returns None if the bundle isn't known.
59 Throws BundleReadException
57 Throws BundleReadException
60 The returned object should be a file object supporting read()
58 The returned object should be a file object supporting read()
61 and close().
59 and close().
62 """
60 """
63
61
64
62
65 class filebundlestore(object):
63 class filebundlestore(object):
66 """bundle store in filesystem
64 """bundle store in filesystem
67
65
68 meant for storing bundles somewhere on disk and on network filesystems
66 meant for storing bundles somewhere on disk and on network filesystems
69 """
67 """
70
68
71 def __init__(self, ui, repo):
69 def __init__(self, ui, repo):
72 self.ui = ui
70 self.ui = ui
73 self.repo = repo
71 self.repo = repo
74 self.storepath = ui.configpath(b'scratchbranch', b'storepath')
72 self.storepath = ui.configpath(b'scratchbranch', b'storepath')
75 if not self.storepath:
73 if not self.storepath:
76 self.storepath = self.repo.vfs.join(
74 self.storepath = self.repo.vfs.join(
77 b"scratchbranches", b"filebundlestore"
75 b"scratchbranches", b"filebundlestore"
78 )
76 )
79 if not os.path.exists(self.storepath):
77 if not os.path.exists(self.storepath):
80 os.makedirs(self.storepath)
78 os.makedirs(self.storepath)
81
79
82 def _dirpath(self, hashvalue):
80 def _dirpath(self, hashvalue):
83 """First two bytes of the hash are the name of the upper
81 """First two bytes of the hash are the name of the upper
84 level directory, next two bytes are the name of the
82 level directory, next two bytes are the name of the
85 next level directory"""
83 next level directory"""
86 return os.path.join(self.storepath, hashvalue[0:2], hashvalue[2:4])
84 return os.path.join(self.storepath, hashvalue[0:2], hashvalue[2:4])
87
85
88 def _filepath(self, filename):
86 def _filepath(self, filename):
89 return os.path.join(self._dirpath(filename), filename)
87 return os.path.join(self._dirpath(filename), filename)
90
88
91 def write(self, data):
89 def write(self, data):
92 filename = node.hex(hashutil.sha1(data).digest())
90 filename = node.hex(hashutil.sha1(data).digest())
93 dirpath = self._dirpath(filename)
91 dirpath = self._dirpath(filename)
94
92
95 if not os.path.exists(dirpath):
93 if not os.path.exists(dirpath):
96 os.makedirs(dirpath)
94 os.makedirs(dirpath)
97
95
98 with open(self._filepath(filename), b'wb') as f:
96 with open(self._filepath(filename), b'wb') as f:
99 f.write(data)
97 f.write(data)
100
98
101 return filename
99 return filename
102
100
103 def read(self, key):
101 def read(self, key):
104 try:
102 try:
105 with open(self._filepath(key), b'rb') as f:
103 with open(self._filepath(key), b'rb') as f:
106 return f.read()
104 return f.read()
107 except IOError:
105 except IOError:
108 return None
106 return None
109
107
110
108
111 class externalbundlestore(abstractbundlestore):
109 class externalbundlestore(abstractbundlestore):
112 def __init__(self, put_binary, put_args, get_binary, get_args):
110 def __init__(self, put_binary, put_args, get_binary, get_args):
113 """
111 """
114 `put_binary` - path to binary file which uploads bundle to external
112 `put_binary` - path to binary file which uploads bundle to external
115 storage and prints key to stdout
113 storage and prints key to stdout
116 `put_args` - format string with additional args to `put_binary`
114 `put_args` - format string with additional args to `put_binary`
117 {filename} replacement field can be used.
115 {filename} replacement field can be used.
118 `get_binary` - path to binary file which accepts filename and key
116 `get_binary` - path to binary file which accepts filename and key
119 (in that order), downloads bundle from store and saves it to file
117 (in that order), downloads bundle from store and saves it to file
120 `get_args` - format string with additional args to `get_binary`.
118 `get_args` - format string with additional args to `get_binary`.
121 {filename} and {handle} replacement field can be used.
119 {filename} and {handle} replacement field can be used.
122 """
120 """
123
121
124 self.put_args = put_args
122 self.put_args = put_args
125 self.get_args = get_args
123 self.get_args = get_args
126 self.put_binary = put_binary
124 self.put_binary = put_binary
127 self.get_binary = get_binary
125 self.get_binary = get_binary
128
126
129 def _call_binary(self, args):
127 def _call_binary(self, args):
130 p = subprocess.Popen(
128 p = subprocess.Popen(
131 pycompat.rapply(procutil.tonativestr, args),
129 pycompat.rapply(procutil.tonativestr, args),
132 stdout=subprocess.PIPE,
130 stdout=subprocess.PIPE,
133 stderr=subprocess.PIPE,
131 stderr=subprocess.PIPE,
134 close_fds=True,
132 close_fds=True,
135 )
133 )
136 stdout, stderr = p.communicate()
134 stdout, stderr = p.communicate()
137 returncode = p.returncode
135 returncode = p.returncode
138 return returncode, stdout, stderr
136 return returncode, stdout, stderr
139
137
140 def write(self, data):
138 def write(self, data):
141 # Won't work on windows because you can't open file second time without
139 # Won't work on windows because you can't open file second time without
142 # closing it
140 # closing it
143 # TODO: rewrite without str.format() and replace NamedTemporaryFile()
141 # TODO: rewrite without str.format() and replace NamedTemporaryFile()
144 # with pycompat.namedtempfile()
142 # with pycompat.namedtempfile()
145 with NamedTemporaryFile() as temp:
143 with pycompat.namedtempfile() as temp:
146 temp.write(data)
144 temp.write(data)
147 temp.flush()
145 temp.flush()
148 temp.seek(0)
146 temp.seek(0)
149 formatted_args = [
147 formatted_args = [
150 arg.format(filename=temp.name) for arg in self.put_args
148 arg.format(filename=temp.name) for arg in self.put_args
151 ]
149 ]
152 returncode, stdout, stderr = self._call_binary(
150 returncode, stdout, stderr = self._call_binary(
153 [self.put_binary] + formatted_args
151 [self.put_binary] + formatted_args
154 )
152 )
155
153
156 if returncode != 0:
154 if returncode != 0:
157 raise BundleWriteException(
155 raise BundleWriteException(
158 b'Failed to upload to external store: %s' % stderr
156 b'Failed to upload to external store: %s' % stderr
159 )
157 )
160 stdout_lines = stdout.splitlines()
158 stdout_lines = stdout.splitlines()
161 if len(stdout_lines) == 1:
159 if len(stdout_lines) == 1:
162 return stdout_lines[0]
160 return stdout_lines[0]
163 else:
161 else:
164 raise BundleWriteException(
162 raise BundleWriteException(
165 b'Bad output from %s: %s' % (self.put_binary, stdout)
163 b'Bad output from %s: %s' % (self.put_binary, stdout)
166 )
164 )
167
165
168 def read(self, handle):
166 def read(self, handle):
169 # Won't work on windows because you can't open file second time without
167 # Won't work on windows because you can't open file second time without
170 # closing it
168 # closing it
171 # TODO: rewrite without str.format() and replace NamedTemporaryFile()
169 # TODO: rewrite without str.format()
172 # with pycompat.namedtempfile()
170 with pycompat.namedtempfile() as temp:
173 with NamedTemporaryFile() as temp:
174 formatted_args = [
171 formatted_args = [
175 arg.format(filename=temp.name, handle=handle)
172 arg.format(filename=temp.name, handle=handle)
176 for arg in self.get_args
173 for arg in self.get_args
177 ]
174 ]
178 returncode, stdout, stderr = self._call_binary(
175 returncode, stdout, stderr = self._call_binary(
179 [self.get_binary] + formatted_args
176 [self.get_binary] + formatted_args
180 )
177 )
181
178
182 if returncode != 0:
179 if returncode != 0:
183 raise BundleReadException(
180 raise BundleReadException(
184 b'Failed to download from external store: %s' % stderr
181 b'Failed to download from external store: %s' % stderr
185 )
182 )
186 return temp.read()
183 return temp.read()
General Comments 0
You need to be logged in to leave comments. Login now