1from os.path import dirname, basename
2
3import os
4
5from omegaml.backends.basedata import BaseDataBackend
6
7try:
8 from smart_open import open
9except:
10 pass
11
12
[docs]
13class PythonRawFileBackend(BaseDataBackend):
14 """
15 OmegaStore backend to support arbitrary files
16 """
17 KIND = 'python.file'
18
[docs]
19 @classmethod
20 def supports(self, obj, name, open_kwargs=None, **kwargs):
21 is_filelike = hasattr(obj, 'read')
22 open_kwargs = dict(open_kwargs or {})
23 if kwargs.get('kind') == self.KIND:
24 is_filelike = self._is_openable(self, obj, **open_kwargs)
25 return is_filelike or self._is_path(self, obj)
26
27 def _is_openable(self, obj, **kwargs):
28 if 'mode' not in 'kwargs':
29 kwargs['mode'] = 'rb'
30 try:
31 with open(obj, **kwargs) as fin:
32 fin.read(1)
33 except:
34 return False
35 return True
36
[docs]
37 def get(self, name, local=None, mode='wb', open_kwargs=None, **kwargs):
38 """
39 get a stored file as a file-like object with binary contents or a local file
40
41 Args:
42 name (str): the name of the file
43 local (str): if set the local path will be created and the file
44 stored there. If local does not have an extension it is assumed
45 to be a directory name, in which case the file is stored as the
46 same name.
47 mode (str): the mode to use on .open() for the local file
48 open_kwargs (dict): the kwargs to use .open() for the local file
49 **kwargs: any kwargs passed to datasets.metadata()
50
51 Returns:
52 the file-like output handler (local is None)
53 the path to the local file (local is given)
54
55 See also:
56 https://docs.python.org/3/glossary.html#term-file-object
57 https://docs.python.org/3/glossary.html#term-binary-file
58 """
59 outf = self.data_store.metadata(name, **kwargs).gridfile
60 if local:
61 is_filename = '.' in basename(local)
62 target_dir = dirname(local) if is_filename else local
63 local = local if is_filename else '{local}/{name}'.format(**locals())
64 os.makedirs(target_dir, exist_ok=True)
65 open_kwargs = open_kwargs or {}
66 with open(local, mode=mode, **open_kwargs) as flocal:
67 flocal.write(outf.read())
68 return local
69 return filelike(outf)
70
[docs]
71 def put(self, obj, name, attributes=None, encoding=None, **kwargs):
72 self.data_store.drop(name, force=True)
73 storekey = self.data_store.object_store_key(name, 'file', hashed=True)
74 gridfile = self._store_to_file(self.data_store, obj, storekey, encoding=encoding)
75 return self.data_store._make_metadata(
76 name=name,
77 prefix=self.data_store.prefix,
78 bucket=self.data_store.bucket,
79 kind=self.KIND,
80 attributes=attributes,
81 gridfile=gridfile).save()
82
83
84def filelike(obj):
85 actual = obj.get()
86 __doc__ = actual.__doc__
87 return actual