1import io
2import logging
3import os
4import zipfile
5from os.path import basename, dirname
6from pathlib import Path
7
8import smart_open
9
10from omegaml.backends.basedata import BaseDataBackend
11from omegaml.util import tryOr
12
13try:
14 from smart_open import open
15except:
16 pass
17
18logger = logging.getLogger(__name__)
19
20
[docs]
21class PythonRawFileBackend(BaseDataBackend):
22 """
23 OmegaStore backend to support arbitrary files
24 """
25
26 KIND = 'python.file'
27
[docs]
28 @classmethod
29 def supports(self, obj, name, open_kwargs=None, **kwargs):
30 is_filelike = hasattr(obj, 'read')
31 open_kwargs = dict(open_kwargs or {})
32 if kwargs.get('kind') == self.KIND:
33 is_filelike |= self._is_openable(self, obj, **open_kwargs)
34 return is_filelike or self._is_path(self, obj)
35
36 def _is_openable(self, obj, **kwargs):
37 if 'mode' not in 'kwargs':
38 kwargs['mode'] = 'rb'
39 # already opened file
40 if isinstance(obj, io.IOBase):
41 return not obj.closed
42 try:
43 with open(obj, **kwargs) as fin:
44 fin.read(1)
45 except:
46 return False
47 return True
48
[docs]
49 def get(
50 self,
51 name,
52 local=None,
53 mode='wb',
54 open_kwargs=None,
55 chunksize=None,
56 uri=None,
57 extract=None,
58 replace=False,
59 **kwargs,
60 ):
61 """
62 get a stored file as a file-like object with binary contents, or copy contents to a local file
63
64 Args:
65 name (str): the name of the file
66 local (str): if set the local path will be created and the file
67 stored there. If local does not have an extension it is assumed
68 to be a directory name, where the file will stored with /path/to/local/name.
69 If the directory does not exist, the directory is created.
70 mode (str): the mode to use on .open() for the local file
71 chunksize (int): optional, the size of chunks to be read, as in
72 open_kwargs (dict): optional, the kwargs to use .open() for the local file
73 uri (str): optional, a uri passed to smart_open for the source location of the file.
74 If not specified, defaults to metadata.uri, which defaults to metadata.gridfile
75 extract (bool): optional, defaults to False. If True and the file is a zipfile, the file will be
76 extracted to the local= path. Defaults to True if local is a directory or if
77 uri exists as a local path; in this case uri is the local path.
78 replace (bool): if True, an existing local file will be overwritten
79 **kwargs: any kwargs passed to datasets.metadata()
80
81 Returns:
82 the file-like output handler (local is None)
83 the path to the local file (local is given)
84
85 See also:
86 https://docs.python.org/3/glossary.html#term-file-object
87 https://docs.python.org/3/glossary.html#term-binary-file
88
89 .. versionchanged:: 0.18.2
90 specify extracxt=True to retrieve the contents of a zipped-up file, see also BackendBaseCommon._store_to_file()
91 """
92 meta = self.data_store.metadata(name, **kwargs)
93 chunksize = chunksize or 1024 * 1024 * 4
94 if uri and Path(uri).parent.exists():
95 local = local or uri
96 else:
97 uri = meta.uri
98 if uri:
99 outf = smart_open.open(uri, mode='rb')
100 else:
101 outf = self.data_store.metadata(name, **kwargs).gridfile
102 if local:
103 # cases:
104 # a. local is a filename (has an extension, is not a directory)
105 # => output contents to file, don't extract
106 # b. stored file is a zip file, local is a directory (filename without extension)
107 # => extract, unless extract is False
108 # => if False, error
109 # c. store file is omegaml-zipped, local is a directory (filename without extension)
110 # => extract, unless extract is False
111 # => if False, error
112 is_filename = not Path(local).is_dir() and (Path(local).is_file() or Path(local).suffix != '')
113 target_dir = dirname(local) if is_filename and not extract else local
114 extract = extract if extract is not None else not is_filename
115 as_file_local = '{local}/{name}'.format(**locals()) if not extract else local
116 local = local if is_filename else as_file_local
117 open_kwargs = open_kwargs or {}
118 is_zipfile = zipfile.is_zipfile(outf)
119 outf.seek(0) # ensure we're back to 0 offset after zipfile.is_zipfile()
120 if is_zipfile:
121 target_dir = Path(target_dir)
122 if target_dir.is_file():
123 target_dir = target_dir.with_suffix('.unzipped')
124 local = target_dir
125 os.makedirs(target_dir, exist_ok=True)
126 with zipfile.ZipFile(outf) as zip:
127 omega_zipped = any(basename(fn) == self._magicszip for fn in zip.namelist())
128 # cases b, c
129 should_extract = extract or (omega_zipped and extract is True)
130 if should_extract:
131 zip.extractall(path=target_dir)
132 # remove directory marker to avoid manifest issues
133 for fn in target_dir.glob(f'**/{self._magicszip}'):
134 fn.unlink(missing_ok=True)
135 if not should_extract:
136 outf.seek(0)
137 if local != uri:
138 self._from_store_to_local(outf, local=local, mode=mode, open_kwargs=open_kwargs)
139 elif Path(local).exists() and not replace:
140 logger.warning(f'{local} exists already, no data written, specify replace=True')
141 else:
142 # case a
143 outf.seek(0)
144 self._from_store_to_local(outf, local=local, mode=mode, open_kwargs=open_kwargs)
145 tryOr(lambda: outf.close(), False)
146 return Path(local)
147 return filelike(outf)
148
[docs]
149 def put(self, obj, name, attributes=None, encoding=None, uri=None, **kwargs):
150 """
151 store the binary contents of a file-like object
152
153 Args:
154 obj (str|Path|filelike): the object to be stored
155 name (str): the name for the object's metadata
156 attributes (dict): optional, metadata attributes
157 encoding (str): optional, a valid encoding, such as utf8
158 uri (str): optional, the local or remote file url compatible with smart_open
159 **kwargs:
160
161 Returns:
162 Metadata
163 """
164 self.data_store.drop(name, force=True)
165 storekey = self.data_store.object_store_key(name, 'file', hashed=True)
166 gridfile = self._store_to_file(self.data_store, obj, storekey, encoding=encoding, uri=uri, **kwargs)
167 return self.data_store._make_metadata(
168 name=name,
169 prefix=self.data_store.prefix,
170 bucket=self.data_store.bucket,
171 kind=self.KIND,
172 attributes=attributes,
173 uri=str(uri or ''),
174 gridfile=gridfile,
175 ).save()
176
177
178def filelike(obj):
179 # convert GridFsProxy to GridOut, a filelike object
180 # -- for actual files, returns just the actual file
181 actual = obj.get() if hasattr(obj, 'get') else obj
182 __doc__ = actual.__doc__
183 return actual