Coverage for python/lsst/daf/persistence/posixStorage.py : 13%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#!/usr/bin/env python
3#
4# LSST Data Management System
5# Copyright 2016 LSST Corporation.
6#
7# This product includes software developed by the
8# LSST Project (http://www.lsst.org/).
9#
10# This program is free software: you can redistribute it and/or modify
11# it under the terms of the GNU General Public License as published by
12# the Free Software Foundation, either version 3 of the License, or
13# (at your option) any later version.
14#
15# This program is distributed in the hope that it will be useful,
16# but WITHOUT ANY WARRANTY; without even the implied warranty of
17# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
18# GNU General Public License for more details.
19#
20# You should have received a copy of the LSST License Statement and
21# the GNU General Public License along with this program. If not,
22# see <http://www.lsstcorp.org/LegalNotices/>.
23#
24import sys
25import pickle
26import importlib
27import os
28import re
29import urllib.parse
30import glob
31import shutil
32import yaml
34from . import (LogicalLocation, Policy,
35 StorageInterface, Storage, ButlerLocation,
36 NoRepositroyAtRoot, RepositoryCfg, doImport)
37from lsst.log import Log
38import lsst.pex.policy as pexPolicy
39from .safeFileIo import SafeFilename, safeMakeDir
42__all__ = ["PosixStorage"]
45class PosixStorage(StorageInterface):
46 """Defines the interface for a storage location on the local filesystem.
48 Parameters
49 ----------
50 uri : string
51 URI or path that is used as the storage location.
52 create : bool
53 If True a new repository will be created at the root location if it
54 does not exist. If False then a new repository will not be created.
56 Raises
57 ------
58 NoRepositroyAtRoot
59 If create is False and a repository does not exist at the root
60 specified by uri then NoRepositroyAtRoot is raised.
61 """
63 def __init__(self, uri, create):
64 self.log = Log.getLogger("daf.persistence.butler")
65 self.root = self._pathFromURI(uri)
66 if self.root and not os.path.exists(self.root):
67 if not create:
68 raise NoRepositroyAtRoot("No repository at {}".format(uri))
69 safeMakeDir(self.root)
71 def __repr__(self):
72 return 'PosixStorage(root=%s)' % self.root
74 @staticmethod
75 def _pathFromURI(uri):
76 """Get the path part of the URI"""
77 return urllib.parse.urlparse(uri).path
79 @staticmethod
80 def relativePath(fromPath, toPath):
81 """Get a relative path from a location to a location.
83 Parameters
84 ----------
85 fromPath : string
86 A path at which to start. It can be a relative path or an
87 absolute path.
88 toPath : string
89 A target location. It can be a relative path or an absolute path.
91 Returns
92 -------
93 string
94 A relative path that describes the path from fromPath to toPath.
95 """
96 fromPath = os.path.realpath(fromPath)
97 return os.path.relpath(toPath, fromPath)
99 @staticmethod
100 def absolutePath(fromPath, relativePath):
101 """Get an absolute path for the path from fromUri to toUri
103 Parameters
104 ----------
105 fromPath : the starting location
106 A location at which to start. It can be a relative path or an
107 absolute path.
108 relativePath : the location relative to fromPath
109 A relative path.
111 Returns
112 -------
113 string
114 Path that is an absolute path representation of fromPath +
115 relativePath, if one exists. If relativePath is absolute or if
116 fromPath is not related to relativePath then relativePath will be
117 returned.
118 """
119 if os.path.isabs(relativePath):
120 return relativePath
121 fromPath = os.path.realpath(fromPath)
122 return os.path.normpath(os.path.join(fromPath, relativePath))
124 @staticmethod
125 def getRepositoryCfg(uri):
126 """Get a persisted RepositoryCfg
128 Parameters
129 ----------
130 uri : URI or path to a RepositoryCfg
131 Description
133 Returns
134 -------
135 A RepositoryCfg instance or None
136 """
137 storage = Storage.makeFromURI(uri)
138 location = ButlerLocation(pythonType=RepositoryCfg,
139 cppType=None,
140 storageName=None,
141 locationList='repositoryCfg.yaml',
142 dataId={},
143 mapper=None,
144 storage=storage,
145 usedDataId=None,
146 datasetType=None)
147 return storage.read(location)
149 @staticmethod
150 def putRepositoryCfg(cfg, loc=None):
151 storage = Storage.makeFromURI(cfg.root if loc is None else loc, create=True)
152 location = ButlerLocation(pythonType=RepositoryCfg,
153 cppType=None,
154 storageName=None,
155 locationList='repositoryCfg.yaml',
156 dataId={},
157 mapper=None,
158 storage=storage,
159 usedDataId=None,
160 datasetType=None)
161 storage.write(location, cfg)
163 @staticmethod
164 def getMapperClass(root):
165 """Get the mapper class associated with a repository root.
167 Supports the legacy _parent symlink search (which was only ever posix-only. This should not be used by
168 new code and repositories; they should use the Repository parentCfg mechanism.
170 Parameters
171 ----------
172 root : string
173 The location of a persisted ReositoryCfg is (new style repos), or
174 the location where a _mapper file is (old style repos).
176 Returns
177 -------
178 A class object or a class instance, depending on the state of the
179 mapper when the repository was created.
180 """
181 if not (root):
182 return None
184 cfg = PosixStorage.getRepositoryCfg(root)
185 if cfg is not None:
186 return cfg.mapper
188 # Find a "_mapper" file containing the mapper class name
189 basePath = root
190 mapperFile = "_mapper"
191 while not os.path.exists(os.path.join(basePath, mapperFile)):
192 # Break abstraction by following _parent links from CameraMapper
193 if os.path.exists(os.path.join(basePath, "_parent")):
194 basePath = os.path.join(basePath, "_parent")
195 else:
196 mapperFile = None
197 break
199 if mapperFile is not None:
200 mapperFile = os.path.join(basePath, mapperFile)
202 # Read the name of the mapper class and instantiate it
203 with open(mapperFile, "r") as f:
204 mapperName = f.readline().strip()
205 components = mapperName.split(".")
206 if len(components) <= 1:
207 raise RuntimeError("Unqualified mapper name %s in %s" %
208 (mapperName, mapperFile))
209 pkg = importlib.import_module(".".join(components[:-1]))
210 return getattr(pkg, components[-1])
212 return None
214 @staticmethod
215 def getParentSymlinkPath(root):
216 """For Butler V1 Repositories only, if a _parent symlink exists, get the location pointed to by the
217 symlink.
219 Parameters
220 ----------
221 root : string
222 A path to the folder on the local filesystem.
224 Returns
225 -------
226 string or None
227 A path to the parent folder indicated by the _parent symlink, or None if there is no _parent
228 symlink at root.
229 """
230 linkpath = os.path.join(root, '_parent')
231 if os.path.exists(linkpath):
232 try:
233 return os.readlink(os.path.join(root, '_parent'))
234 except OSError:
235 # some of the unit tests rely on a folder called _parent instead of a symlink to aother
236 # location. Allow that; return the path of that folder.
237 return os.path.join(root, '_parent')
238 return None
240 def write(self, butlerLocation, obj):
241 """Writes an object to a location and persistence format specified by
242 ButlerLocation
244 Parameters
245 ----------
246 butlerLocation : ButlerLocation
247 The location & formatting for the object to be written.
248 obj : object instance
249 The object to be written.
250 """
251 self.log.debug("Put location=%s obj=%s", butlerLocation, obj)
253 writeFormatter = self.getWriteFormatter(butlerLocation.getStorageName())
254 if not writeFormatter:
255 writeFormatter = self.getWriteFormatter(butlerLocation.getPythonType())
256 if writeFormatter:
257 writeFormatter(butlerLocation, obj)
258 return
260 raise(RuntimeError("No formatter for location:{}".format(butlerLocation)))
262 def read(self, butlerLocation):
263 """Read from a butlerLocation.
265 Parameters
266 ----------
267 butlerLocation : ButlerLocation
268 The location & formatting for the object(s) to be read.
270 Returns
271 -------
272 A list of objects as described by the butler location. One item for
273 each location in butlerLocation.getLocations()
274 """
275 readFormatter = self.getReadFormatter(butlerLocation.getStorageName())
276 if not readFormatter:
277 readFormatter = self.getReadFormatter(butlerLocation.getPythonType())
278 if readFormatter:
279 return readFormatter(butlerLocation)
281 raise(RuntimeError("No formatter for location:{}".format(butlerLocation)))
283 def butlerLocationExists(self, location):
284 """Implementation of PosixStorage.exists for ButlerLocation objects.
285 """
286 storageName = location.getStorageName()
287 if storageName not in ('FitsStorage', 'PafStorage',
288 'PickleStorage', 'ConfigStorage', 'FitsCatalogStorage',
289 'YamlStorage', 'ParquetStorage', 'MatplotlibStorage'):
290 self.log.warn("butlerLocationExists for non-supported storage %s" % location)
291 return False
292 for locationString in location.getLocations():
293 logLoc = LogicalLocation(locationString, location.getAdditionalData()).locString()
294 obj = self.instanceSearch(path=logLoc)
295 if obj:
296 return True
297 return False
299 def exists(self, location):
300 """Check if location exists.
302 Parameters
303 ----------
304 location : ButlerLocation or string
305 A a string or a ButlerLocation that describes the location of an
306 object in this storage.
308 Returns
309 -------
310 bool
311 True if exists, else False.
312 """
313 if isinstance(location, ButlerLocation):
314 return self.butlerLocationExists(location)
316 obj = self.instanceSearch(path=location)
317 return bool(obj)
319 def locationWithRoot(self, location):
320 """Get the full path to the location.
322 :param location:
323 :return:
324 """
325 return os.path.join(self.root, location)
327 @staticmethod
328 def v1RepoExists(root):
329 """Test if a Version 1 Repository exists.
331 Version 1 Repositories only exist in posix storages, do not have a
332 RepositoryCfg file, and contain either a registry.sqlite3 file, a
333 _mapper file, or a _parent link.
335 Parameters
336 ----------
337 root : string
338 A path to a folder on the local filesystem.
340 Returns
341 -------
342 bool
343 True if the repository at root exists, else False.
344 """
345 return os.path.exists(root) and (
346 os.path.exists(os.path.join(root, "registry.sqlite3"))
347 or os.path.exists(os.path.join(root, "_mapper"))
348 or os.path.exists(os.path.join(root, "_parent"))
349 )
351 def copyFile(self, fromLocation, toLocation):
352 """Copy a file from one location to another on the local filesystem.
354 Parameters
355 ----------
356 fromLocation : path
357 Path and name of existing file.
358 toLocation : path
359 Path and name of new file.
361 Returns
362 -------
363 None
364 """
365 shutil.copy(os.path.join(self.root, fromLocation), os.path.join(self.root, toLocation))
367 def getLocalFile(self, path):
368 """Get a handle to a local copy of the file, downloading it to a
369 temporary if needed.
371 Parameters
372 ----------
373 A path the the file in storage, relative to root.
375 Returns
376 -------
377 A handle to a local copy of the file. If storage is remote it will be
378 a temporary file. If storage is local it may be the original file or
379 a temporary file. The file name can be gotten via the 'name' property
380 of the returned object.
381 """
382 p = os.path.join(self.root, path)
383 try:
384 return open(p)
385 except IOError as e:
386 if e.errno == 2: # 'No such file or directory'
387 return None
388 else:
389 raise e
391 def instanceSearch(self, path):
392 """Search for the given path in this storage instance.
394 If the path contains an HDU indicator (a number in brackets before the
395 dot, e.g. 'foo.fits[1]', this will be stripped when searching and so
396 will match filenames without the HDU indicator, e.g. 'foo.fits'. The
397 path returned WILL contain the indicator though, e.g. ['foo.fits[1]'].
399 Parameters
400 ----------
401 path : string
402 A filename (and optionally prefix path) to search for within root.
404 Returns
405 -------
406 string or None
407 The location that was found, or None if no location was found.
408 """
409 return self.search(self.root, path)
411 @staticmethod
412 def search(root, path, searchParents=False):
413 """Look for the given path in the current root.
415 Also supports searching for the path in Butler v1 repositories by
416 following the Butler v1 _parent symlink
418 If the path contains an HDU indicator (a number in brackets, e.g.
419 'foo.fits[1]', this will be stripped when searching and so
420 will match filenames without the HDU indicator, e.g. 'foo.fits'. The
421 path returned WILL contain the indicator though, e.g. ['foo.fits[1]'].
423 Parameters
424 ----------
425 root : string
426 The path to the root directory.
427 path : string
428 The path to the file within the root directory.
429 searchParents : bool, optional
430 For Butler v1 repositories only, if true and a _parent symlink
431 exists, then the directory at _parent will be searched if the file
432 is not found in the root repository. Will continue searching the
433 parent of the parent until the file is found or no additional
434 parent exists.
436 Returns
437 -------
438 string or None
439 The location that was found, or None if no location was found.
440 """
441 # Separate path into a root-equivalent prefix (in dir) and the rest
442 # (left in path)
443 rootDir = root
444 # First remove trailing slashes (#2527)
445 while len(rootDir) > 1 and rootDir[-1] == '/':
446 rootDir = rootDir[:-1]
448 if not path.startswith('/'):
449 # Most common case is a relative path from a template
450 pathPrefix = None
451 elif path.startswith(rootDir + "/"):
452 # Common case; we have the same root prefix string
453 path = path[len(rootDir + '/'):]
454 pathPrefix = rootDir
455 elif rootDir == "/" and path.startswith("/"):
456 path = path[1:]
457 pathPrefix = None
458 else:
459 # Search for prefix that is the same as root
460 pathPrefix = os.path.dirname(path)
461 while pathPrefix != "" and pathPrefix != "/":
462 if os.path.realpath(pathPrefix) == os.path.realpath(root):
463 break
464 pathPrefix = os.path.dirname(pathPrefix)
465 if pathPrefix == "/":
466 path = path[1:]
467 elif pathPrefix != "":
468 path = path[len(pathPrefix)+1:]
470 # Now search for the path in the root or its parents
471 # Strip off any cfitsio bracketed extension if present
472 strippedPath = path
473 pathStripped = None
474 firstBracket = path.find("[")
475 if firstBracket != -1:
476 strippedPath = path[:firstBracket]
477 pathStripped = path[firstBracket:]
479 dir = rootDir
480 while True:
481 paths = glob.glob(os.path.join(dir, strippedPath))
482 if len(paths) > 0:
483 if pathPrefix != rootDir:
484 paths = [p[len(rootDir+'/'):] for p in paths]
485 if pathStripped is not None:
486 paths = [p + pathStripped for p in paths]
487 return paths
488 if searchParents:
489 dir = os.path.join(dir, "_parent")
490 if not os.path.exists(dir):
491 return None
492 else:
493 return None
495 @staticmethod
496 def storageExists(uri):
497 """Ask if a storage at the location described by uri exists
499 Parameters
500 ----------
501 root : string
502 URI to the the root location of the storage
504 Returns
505 -------
506 bool
507 True if the storage exists, false if not
508 """
509 return os.path.exists(PosixStorage._pathFromURI(uri))
512def readConfigStorage(butlerLocation):
513 """Read an lsst.pex.config.Config from a butlerLocation.
515 Parameters
516 ----------
517 butlerLocation : ButlerLocation
518 The location for the object(s) to be read.
520 Returns
521 -------
522 A list of objects as described by the butler location. One item for
523 each location in butlerLocation.getLocations()
524 """
525 results = []
526 for locationString in butlerLocation.getLocations():
527 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString)
528 logLoc = LogicalLocation(locStringWithRoot, butlerLocation.getAdditionalData())
529 if not os.path.exists(logLoc.locString()):
530 raise RuntimeError("No such config file: " + logLoc.locString())
531 pythonType = butlerLocation.getPythonType()
532 if pythonType is not None:
533 if isinstance(pythonType, str):
534 pythonType = doImport(pythonType)
535 finalItem = pythonType()
536 finalItem.load(logLoc.locString())
537 results.append(finalItem)
538 return results
541def writeConfigStorage(butlerLocation, obj):
542 """Writes an lsst.pex.config.Config object to a location specified by
543 ButlerLocation.
545 Parameters
546 ----------
547 butlerLocation : ButlerLocation
548 The location for the object to be written.
549 obj : object instance
550 The object to be written.
551 """
552 filename = os.path.join(butlerLocation.getStorage().root, butlerLocation.getLocations()[0])
553 with SafeFilename(filename) as locationString:
554 logLoc = LogicalLocation(locationString, butlerLocation.getAdditionalData())
555 obj.save(logLoc.locString())
558def readFitsStorage(butlerLocation):
559 """Read objects from a FITS file specified by ButlerLocation.
561 The object is read using class or static method
562 ``readFitsWithOptions(path, options)``, if it exists, else
563 ``readFits(path)``. The ``options`` argument is the data returned by
564 ``butlerLocation.getAdditionalData()``.
566 Parameters
567 ----------
568 butlerLocation : ButlerLocation
569 The location for the object(s) to be read.
571 Returns
572 -------
573 A list of objects as described by the butler location. One item for
574 each location in butlerLocation.getLocations()
575 """
576 pythonType = butlerLocation.getPythonType()
577 if pythonType is not None:
578 if isinstance(pythonType, str):
579 pythonType = doImport(pythonType)
580 supportsOptions = hasattr(pythonType, "readFitsWithOptions")
581 if not supportsOptions:
582 from lsst.daf.base import PropertySet, PropertyList
583 if issubclass(pythonType, (PropertySet, PropertyList)):
584 from lsst.afw.image import readMetadata
585 reader = readMetadata
586 else:
587 reader = pythonType.readFits
588 results = []
589 additionalData = butlerLocation.getAdditionalData()
590 for locationString in butlerLocation.getLocations():
591 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString)
592 logLoc = LogicalLocation(locStringWithRoot, additionalData)
593 # test for existence of file, ignoring trailing [...]
594 # because that can specify the HDU or other information
595 filePath = re.sub(r"(\.fits(.[a-zA-Z0-9]+)?)(\[.+\])$", r"\1", logLoc.locString())
596 if not os.path.exists(filePath):
597 raise RuntimeError("No such FITS file: " + logLoc.locString())
598 if supportsOptions:
599 finalItem = pythonType.readFitsWithOptions(logLoc.locString(), options=additionalData)
600 else:
601 fileName = logLoc.locString()
602 mat = re.search(r"^(.*)\[(\d+)\]$", fileName)
604 if mat and reader == readMetadata: # readMetadata() only understands the hdu argument, not [hdu]
605 fileName = mat.group(1)
606 hdu = int(mat.group(2))
608 finalItem = reader(fileName, hdu=hdu)
609 else:
610 finalItem = reader(fileName)
611 results.append(finalItem)
612 return results
615def writeFitsStorage(butlerLocation, obj):
616 """Writes an object to a FITS file specified by ButlerLocation.
618 The object is written using method
619 ``writeFitsWithOptions(path, options)``, if it exists, else
620 ``writeFits(path)``. The ``options`` argument is the data returned by
621 ``butlerLocation.getAdditionalData()``.
623 Parameters
624 ----------
625 butlerLocation : ButlerLocation
626 The location for the object to be written.
627 obj : object instance
628 The object to be written.
629 """
630 supportsOptions = hasattr(obj, "writeFitsWithOptions")
631 additionalData = butlerLocation.getAdditionalData()
632 locations = butlerLocation.getLocations()
633 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString:
634 logLoc = LogicalLocation(locationString, additionalData)
635 if supportsOptions:
636 obj.writeFitsWithOptions(logLoc.locString(), options=additionalData)
637 else:
638 obj.writeFits(logLoc.locString())
641def readParquetStorage(butlerLocation):
642 """Read a catalog from a Parquet file specified by ButlerLocation.
644 The object returned by this is expected to be a subtype
645 of `ParquetTable`, which is a thin wrapper to `pyarrow.ParquetFile`
646 that allows for lazy loading of the data.
648 Parameters
649 ----------
650 butlerLocation : ButlerLocation
651 The location for the object(s) to be read.
653 Returns
654 -------
655 A list of objects as described by the butler location. One item for
656 each location in butlerLocation.getLocations()
657 """
658 results = []
659 additionalData = butlerLocation.getAdditionalData()
661 for locationString in butlerLocation.getLocations():
662 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString)
663 logLoc = LogicalLocation(locStringWithRoot, additionalData)
664 if not os.path.exists(logLoc.locString()):
665 raise RuntimeError("No such parquet file: " + logLoc.locString())
667 pythonType = butlerLocation.getPythonType()
668 if pythonType is not None:
669 if isinstance(pythonType, str):
670 pythonType = doImport(pythonType)
672 filename = logLoc.locString()
674 # pythonType will be ParquetTable (or perhaps MultilevelParquetTable)
675 # filename should be the first kwarg, but being explicit here.
676 results.append(pythonType(filename=filename))
678 return results
681def writeParquetStorage(butlerLocation, obj):
682 """Writes pandas dataframe to parquet file.
684 Parameters
685 ----------
686 butlerLocation : ButlerLocation
687 The location for the object(s) to be read.
688 obj : `lsst.qa.explorer.parquetTable.ParquetTable`
689 Wrapped DataFrame to write.
691 """
692 additionalData = butlerLocation.getAdditionalData()
693 locations = butlerLocation.getLocations()
694 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString:
695 logLoc = LogicalLocation(locationString, additionalData)
696 filename = logLoc.locString()
697 obj.write(filename)
700def writeYamlStorage(butlerLocation, obj):
701 """Writes an object to a YAML file specified by ButlerLocation.
703 Parameters
704 ----------
705 butlerLocation : ButlerLocation
706 The location for the object to be written.
707 obj : object instance
708 The object to be written.
709 """
710 additionalData = butlerLocation.getAdditionalData()
711 locations = butlerLocation.getLocations()
712 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString:
713 logLoc = LogicalLocation(locationString, additionalData)
714 with open(logLoc.locString(), "w") as outfile:
715 yaml.dump(obj, outfile)
718def readPickleStorage(butlerLocation):
719 """Read an object from a pickle file specified by ButlerLocation.
721 Parameters
722 ----------
723 butlerLocation : ButlerLocation
724 The location for the object(s) to be read.
726 Returns
727 -------
728 A list of objects as described by the butler location. One item for
729 each location in butlerLocation.getLocations()
730 """
731 # Create a list of Storages for the item.
732 results = []
733 additionalData = butlerLocation.getAdditionalData()
734 for locationString in butlerLocation.getLocations():
735 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString)
736 logLoc = LogicalLocation(locStringWithRoot, additionalData)
737 if not os.path.exists(logLoc.locString()):
738 raise RuntimeError("No such pickle file: " + logLoc.locString())
739 with open(logLoc.locString(), "rb") as infile:
740 # py3: We have to specify encoding since some files were written
741 # by python2, and 'latin1' manages that conversion safely. See:
742 # http://stackoverflow.com/questions/28218466/unpickling-a-python-2-object-with-python-3/28218598#28218598
743 if sys.version_info.major >= 3:
744 finalItem = pickle.load(infile, encoding="latin1")
745 else:
746 finalItem = pickle.load(infile)
747 results.append(finalItem)
748 return results
751def writePickleStorage(butlerLocation, obj):
752 """Writes an object to a pickle file specified by ButlerLocation.
754 Parameters
755 ----------
756 butlerLocation : ButlerLocation
757 The location for the object to be written.
758 obj : object instance
759 The object to be written.
760 """
761 additionalData = butlerLocation.getAdditionalData()
762 locations = butlerLocation.getLocations()
763 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString:
764 logLoc = LogicalLocation(locationString, additionalData)
765 with open(logLoc.locString(), "wb") as outfile:
766 pickle.dump(obj, outfile, pickle.HIGHEST_PROTOCOL)
769def readFitsCatalogStorage(butlerLocation):
770 """Read a catalog from a FITS table specified by ButlerLocation.
772 Parameters
773 ----------
774 butlerLocation : ButlerLocation
775 The location for the object(s) to be read.
777 Returns
778 -------
779 A list of objects as described by the butler location. One item for
780 each location in butlerLocation.getLocations()
781 """
782 pythonType = butlerLocation.getPythonType()
783 if pythonType is not None:
784 if isinstance(pythonType, str):
785 pythonType = doImport(pythonType)
786 results = []
787 additionalData = butlerLocation.getAdditionalData()
788 for locationString in butlerLocation.getLocations():
789 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString)
790 logLoc = LogicalLocation(locStringWithRoot, additionalData)
791 if not os.path.exists(logLoc.locString()):
792 raise RuntimeError("No such FITS catalog file: " + logLoc.locString())
793 kwds = {}
794 if additionalData.exists("hdu"):
795 kwds["hdu"] = additionalData.getInt("hdu")
796 if additionalData.exists("flags"):
797 kwds["flags"] = additionalData.getInt("flags")
798 finalItem = pythonType.readFits(logLoc.locString(), **kwds)
799 results.append(finalItem)
800 return results
803def writeFitsCatalogStorage(butlerLocation, obj):
804 """Writes a catalog to a FITS table specified by ButlerLocation.
806 Parameters
807 ----------
808 butlerLocation : ButlerLocation
809 The location for the object to be written.
810 obj : object instance
811 The object to be written.
812 """
813 additionalData = butlerLocation.getAdditionalData()
814 locations = butlerLocation.getLocations()
815 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString:
816 logLoc = LogicalLocation(locationString, additionalData)
817 if additionalData.exists("flags"):
818 kwds = dict(flags=additionalData.getInt("flags"))
819 else:
820 kwds = {}
821 obj.writeFits(logLoc.locString(), **kwds)
824def readMatplotlibStorage(butlerLocation):
825 """Read from a butlerLocation (always fails for this storage type).
827 Parameters
828 ----------
829 butlerLocation : ButlerLocation
830 The location for the object(s) to be read.
832 Returns
833 -------
834 A list of objects as described by the butler location. One item for
835 each location in butlerLocation.getLocations()
836 """
837 raise NotImplementedError("Figures saved with MatplotlibStorage cannot be retreived using the Butler.")
840def writeMatplotlibStorage(butlerLocation, obj):
841 """Writes a matplotlib.figure.Figure to a location, using the template's
842 filename suffix to infer the file format.
844 Parameters
845 ----------
846 butlerLocation : ButlerLocation
847 The location for the object to be written.
848 obj : matplotlib.figure.Figure
849 The object to be written.
850 """
851 additionalData = butlerLocation.getAdditionalData()
852 locations = butlerLocation.getLocations()
853 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString:
854 logLoc = LogicalLocation(locationString, additionalData)
855 # SafeFilename appends a random suffix, which corrupts the extension
856 # matplotlib uses to guess the file format.
857 # Instead, we extract the extension from the original location
858 # and pass that as the format directly.
859 _, ext = os.path.splitext(locations[0])
860 if ext:
861 ext = ext[1:] # strip off leading '.'
862 else:
863 # If there is no extension, we let matplotlib fall back to its
864 # default.
865 ext = None
866 obj.savefig(logLoc.locString(), format=ext)
869def readPafStorage(butlerLocation):
870 """Read a policy from a PAF file specified by a ButlerLocation.
872 Parameters
873 ----------
874 butlerLocation : ButlerLocation
875 The location for the object(s) to be read.
877 Returns
878 -------
879 A list of objects as described by the butler location. One item for
880 each location in butlerLocation.getLocations()
881 """
882 results = []
883 for locationString in butlerLocation.getLocations():
884 logLoc = LogicalLocation(butlerLocation.getStorage().locationWithRoot(locationString),
885 butlerLocation.getAdditionalData())
886 finalItem = pexPolicy.Policy.createPolicy(logLoc.locString())
887 results.append(finalItem)
888 return results
891def readYamlStorage(butlerLocation):
892 """Read an object from a YAML file specified by a butlerLocation.
894 Parameters
895 ----------
896 butlerLocation : ButlerLocation
897 The location for the object(s) to be read.
899 Returns
900 -------
901 A list of objects as described by the butler location. One item for
902 each location in butlerLocation.getLocations()
903 """
904 results = []
905 for locationString in butlerLocation.getLocations():
906 logLoc = LogicalLocation(butlerLocation.getStorage().locationWithRoot(locationString),
907 butlerLocation.getAdditionalData())
908 if not os.path.exists(logLoc.locString()):
909 raise RuntimeError("No such YAML file: " + logLoc.locString())
910 # Butler Gen2 repository configurations are handled specially
911 if butlerLocation.pythonType == 'lsst.daf.persistence.RepositoryCfg':
912 finalItem = Policy(filePath=logLoc.locString())
913 else:
914 try:
915 # PyYAML >=5.1 prefers a different loader
916 loader = yaml.FullLoader
917 except AttributeError:
918 loader = yaml.Loader
919 with open(logLoc.locString(), "rb") as infile:
920 finalItem = yaml.load(infile, Loader=loader)
921 results.append(finalItem)
922 return results
925PosixStorage.registerFormatters("FitsStorage", readFitsStorage, writeFitsStorage)
926PosixStorage.registerFormatters("ParquetStorage", readParquetStorage, writeParquetStorage)
927PosixStorage.registerFormatters("ConfigStorage", readConfigStorage, writeConfigStorage)
928PosixStorage.registerFormatters("PickleStorage", readPickleStorage, writePickleStorage)
929PosixStorage.registerFormatters("FitsCatalogStorage", readFitsCatalogStorage, writeFitsCatalogStorage)
930PosixStorage.registerFormatters("MatplotlibStorage", readMatplotlibStorage, writeMatplotlibStorage)
931PosixStorage.registerFormatters("PafStorage", readFormatter=readPafStorage)
932PosixStorage.registerFormatters("YamlStorage", readYamlStorage, writeYamlStorage)
934Storage.registerStorageClass(scheme='', cls=PosixStorage)
935Storage.registerStorageClass(scheme='file', cls=PosixStorage)