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

353 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-14 03:44 -0800

1#!/usr/bin/env python 

2 

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 

33 

34from . import (LogicalLocation, Policy, 

35 StorageInterface, Storage, ButlerLocation, 

36 NoRepositroyAtRoot, RepositoryCfg, doImport) 

37from lsst.log import Log 

38import lsst.pex.config as pexConfig 

39from .safeFileIo import SafeFilename, safeMakeDir 

40 

41 

42__all__ = ["PosixStorage"] 

43 

44 

45class PosixStorage(StorageInterface): 

46 """Defines the interface for a storage location on the local filesystem. 

47 

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. 

55 

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 """ 

62 

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) 

70 

71 def __repr__(self): 

72 return 'PosixStorage(root=%s)' % self.root 

73 

74 @staticmethod 

75 def _pathFromURI(uri): 

76 """Get the path part of the URI""" 

77 return urllib.parse.urlparse(uri).path 

78 

79 @staticmethod 

80 def relativePath(fromPath, toPath): 

81 """Get a relative path from a location to a location. 

82 

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. 

90 

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) 

98 

99 @staticmethod 

100 def absolutePath(fromPath, relativePath): 

101 """Get an absolute path for the path from fromUri to toUri 

102 

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. 

110 

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)) 

123 

124 @staticmethod 

125 def getRepositoryCfg(uri): 

126 """Get a persisted RepositoryCfg 

127 

128 Parameters 

129 ---------- 

130 uri : URI or path to a RepositoryCfg 

131 Description 

132 

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) 

148 

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) 

162 

163 @staticmethod 

164 def getMapperClass(root): 

165 """Get the mapper class associated with a repository root. 

166 

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. 

169 

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). 

175 

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 

183 

184 cfg = PosixStorage.getRepositoryCfg(root) 

185 if cfg is not None: 

186 return cfg.mapper 

187 

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 

198 

199 if mapperFile is not None: 

200 mapperFile = os.path.join(basePath, mapperFile) 

201 

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]) 

211 

212 return None 

213 

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. 

218 

219 Parameters 

220 ---------- 

221 root : string 

222 A path to the folder on the local filesystem. 

223 

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 

239 

240 def write(self, butlerLocation, obj): 

241 """Writes an object to a location and persistence format specified by 

242 ButlerLocation 

243 

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) 

252 

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 

259 

260 raise(RuntimeError("No formatter for location:{}".format(butlerLocation))) 

261 

262 def read(self, butlerLocation): 

263 """Read from a butlerLocation. 

264 

265 Parameters 

266 ---------- 

267 butlerLocation : ButlerLocation 

268 The location & formatting for the object(s) to be read. 

269 

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) 

280 

281 raise(RuntimeError("No formatter for location:{}".format(butlerLocation))) 

282 

283 def butlerLocationExists(self, location): 

284 """Implementation of PosixStorage.exists for ButlerLocation objects. 

285 """ 

286 storageName = location.getStorageName() 

287 if storageName not in ('FitsStorage', 

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 

298 

299 def exists(self, location): 

300 """Check if location exists. 

301 

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. 

307 

308 Returns 

309 ------- 

310 bool 

311 True if exists, else False. 

312 """ 

313 if isinstance(location, ButlerLocation): 

314 return self.butlerLocationExists(location) 

315 

316 obj = self.instanceSearch(path=location) 

317 return bool(obj) 

318 

319 def locationWithRoot(self, location): 

320 """Get the full path to the location. 

321 

322 :param location: 

323 :return: 

324 """ 

325 return os.path.join(self.root, location) 

326 

327 @staticmethod 

328 def v1RepoExists(root): 

329 """Test if a Version 1 Repository exists. 

330 

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. 

334 

335 Parameters 

336 ---------- 

337 root : string 

338 A path to a folder on the local filesystem. 

339 

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 ) 

350 

351 def copyFile(self, fromLocation, toLocation): 

352 """Copy a file from one location to another on the local filesystem. 

353 

354 Parameters 

355 ---------- 

356 fromLocation : path 

357 Path and name of existing file. 

358 toLocation : path 

359 Path and name of new file. 

360 

361 Returns 

362 ------- 

363 None 

364 """ 

365 shutil.copy(os.path.join(self.root, fromLocation), os.path.join(self.root, toLocation)) 

366 

367 def getLocalFile(self, path): 

368 """Get a handle to a local copy of the file, downloading it to a 

369 temporary if needed. 

370 

371 Parameters 

372 ---------- 

373 A path the the file in storage, relative to root. 

374 

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 

390 

391 def instanceSearch(self, path): 

392 """Search for the given path in this storage instance. 

393 

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]']. 

398 

399 Parameters 

400 ---------- 

401 path : string 

402 A filename (and optionally prefix path) to search for within root. 

403 

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) 

410 

411 @staticmethod 

412 def search(root, path, searchParents=False): 

413 """Look for the given path in the current root. 

414 

415 Also supports searching for the path in Butler v1 repositories by 

416 following the Butler v1 _parent symlink 

417 

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]']. 

422 

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. 

435 

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] 

447 

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:] 

469 

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:] 

478 

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 

494 

495 @staticmethod 

496 def storageExists(uri): 

497 """Ask if a storage at the location described by uri exists 

498 

499 Parameters 

500 ---------- 

501 root : string 

502 URI to the the root location of the storage 

503 

504 Returns 

505 ------- 

506 bool 

507 True if the storage exists, false if not 

508 """ 

509 return os.path.exists(PosixStorage._pathFromURI(uri)) 

510 

511 

512def readConfigStorage(butlerLocation): 

513 """Read an lsst.pex.config.Config from a butlerLocation. 

514 

515 Parameters 

516 ---------- 

517 butlerLocation : ButlerLocation 

518 The location for the object(s) to be read. 

519 

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 

532 # Automatically determine the Config class from the serialized form 

533 with open(logLoc.locString(), "r") as fd: 

534 config_py = fd.read() 

535 config = pexConfig.Config._fromPython(config_py) 

536 

537 pythonType = butlerLocation.getPythonType() 

538 if pythonType is not None: 

539 if isinstance(pythonType, str): 

540 pythonType = doImport(pythonType) 

541 if not isinstance(config, pythonType): 

542 raise TypeError(f"Unexpected type of config: {type(config)}, expected {pythonType}") 

543 

544 results.append(config) 

545 return results 

546 

547 

548def writeConfigStorage(butlerLocation, obj): 

549 """Writes an lsst.pex.config.Config object to a location specified by 

550 ButlerLocation. 

551 

552 Parameters 

553 ---------- 

554 butlerLocation : ButlerLocation 

555 The location for the object to be written. 

556 obj : object instance 

557 The object to be written. 

558 """ 

559 filename = os.path.join(butlerLocation.getStorage().root, butlerLocation.getLocations()[0]) 

560 with SafeFilename(filename) as locationString: 

561 logLoc = LogicalLocation(locationString, butlerLocation.getAdditionalData()) 

562 obj.save(logLoc.locString()) 

563 

564 

565def readFitsStorage(butlerLocation): 

566 """Read objects from a FITS file specified by ButlerLocation. 

567 

568 The object is read using class or static method 

569 ``readFitsWithOptions(path, options)``, if it exists, else 

570 ``readFits(path)``. The ``options`` argument is the data returned by 

571 ``butlerLocation.getAdditionalData()``. 

572 

573 Parameters 

574 ---------- 

575 butlerLocation : ButlerLocation 

576 The location for the object(s) to be read. 

577 

578 Returns 

579 ------- 

580 A list of objects as described by the butler location. One item for 

581 each location in butlerLocation.getLocations() 

582 """ 

583 pythonType = butlerLocation.getPythonType() 

584 if pythonType is not None: 

585 if isinstance(pythonType, str): 

586 pythonType = doImport(pythonType) 

587 supportsOptions = hasattr(pythonType, "readFitsWithOptions") 

588 if not supportsOptions: 

589 from lsst.daf.base import PropertySet, PropertyList 

590 if issubclass(pythonType, (PropertySet, PropertyList)): 

591 from lsst.afw.fits import readMetadata 

592 reader = readMetadata 

593 else: 

594 reader = pythonType.readFits 

595 results = [] 

596 additionalData = butlerLocation.getAdditionalData() 

597 for locationString in butlerLocation.getLocations(): 

598 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString) 

599 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

600 # test for existence of file, ignoring trailing [...] 

601 # because that can specify the HDU or other information 

602 filePath = re.sub(r"(\.fits(.[a-zA-Z0-9]+)?)(\[.+\])$", r"\1", logLoc.locString()) 

603 if not os.path.exists(filePath): 

604 raise RuntimeError("No such FITS file: " + logLoc.locString()) 

605 if supportsOptions: 

606 finalItem = pythonType.readFitsWithOptions(logLoc.locString(), options=additionalData) 

607 else: 

608 fileName = logLoc.locString() 

609 mat = re.search(r"^(.*)\[(\d+)\]$", fileName) 

610 

611 if mat and reader == readMetadata: # readMetadata() only understands the hdu argument, not [hdu] 

612 fileName = mat.group(1) 

613 hdu = int(mat.group(2)) 

614 

615 finalItem = reader(fileName, hdu=hdu) 

616 else: 

617 finalItem = reader(fileName) 

618 results.append(finalItem) 

619 return results 

620 

621 

622def writeFitsStorage(butlerLocation, obj): 

623 """Writes an object to a FITS file specified by ButlerLocation. 

624 

625 The object is written using method 

626 ``writeFitsWithOptions(path, options)``, if it exists, else 

627 ``writeFits(path)``. The ``options`` argument is the data returned by 

628 ``butlerLocation.getAdditionalData()``. 

629 

630 Parameters 

631 ---------- 

632 butlerLocation : ButlerLocation 

633 The location for the object to be written. 

634 obj : object instance 

635 The object to be written. 

636 """ 

637 supportsOptions = hasattr(obj, "writeFitsWithOptions") 

638 additionalData = butlerLocation.getAdditionalData() 

639 locations = butlerLocation.getLocations() 

640 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString: 

641 logLoc = LogicalLocation(locationString, additionalData) 

642 if supportsOptions: 

643 obj.writeFitsWithOptions(logLoc.locString(), options=additionalData) 

644 else: 

645 obj.writeFits(logLoc.locString()) 

646 

647 

648def readParquetStorage(butlerLocation): 

649 """Read a catalog from a Parquet file specified by ButlerLocation. 

650 

651 The object returned by this is expected to be a subtype 

652 of `ParquetTable`, which is a thin wrapper to `pyarrow.ParquetFile` 

653 that allows for lazy loading of the data. 

654 

655 Parameters 

656 ---------- 

657 butlerLocation : ButlerLocation 

658 The location for the object(s) to be read. 

659 

660 Returns 

661 ------- 

662 A list of objects as described by the butler location. One item for 

663 each location in butlerLocation.getLocations() 

664 """ 

665 results = [] 

666 additionalData = butlerLocation.getAdditionalData() 

667 

668 for locationString in butlerLocation.getLocations(): 

669 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString) 

670 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

671 if not os.path.exists(logLoc.locString()): 

672 raise RuntimeError("No such parquet file: " + logLoc.locString()) 

673 

674 pythonType = butlerLocation.getPythonType() 

675 if pythonType is not None: 

676 if isinstance(pythonType, str): 

677 pythonType = doImport(pythonType) 

678 

679 filename = logLoc.locString() 

680 

681 # pythonType will be ParquetTable (or perhaps MultilevelParquetTable) 

682 # filename should be the first kwarg, but being explicit here. 

683 results.append(pythonType(filename=filename)) 

684 

685 return results 

686 

687 

688def writeParquetStorage(butlerLocation, obj): 

689 """Writes pandas dataframe to parquet file. 

690 

691 Parameters 

692 ---------- 

693 butlerLocation : ButlerLocation 

694 The location for the object(s) to be read. 

695 obj : `lsst.qa.explorer.parquetTable.ParquetTable` 

696 Wrapped DataFrame to write. 

697 

698 """ 

699 additionalData = butlerLocation.getAdditionalData() 

700 locations = butlerLocation.getLocations() 

701 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString: 

702 logLoc = LogicalLocation(locationString, additionalData) 

703 filename = logLoc.locString() 

704 obj.write(filename) 

705 

706 

707def writeYamlStorage(butlerLocation, obj): 

708 """Writes an object to a YAML file specified by ButlerLocation. 

709 

710 Parameters 

711 ---------- 

712 butlerLocation : ButlerLocation 

713 The location for the object to be written. 

714 obj : object instance 

715 The object to be written. 

716 """ 

717 additionalData = butlerLocation.getAdditionalData() 

718 locations = butlerLocation.getLocations() 

719 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString: 

720 logLoc = LogicalLocation(locationString, additionalData) 

721 with open(logLoc.locString(), "w") as outfile: 

722 yaml.dump(obj, outfile) 

723 

724 

725def readPickleStorage(butlerLocation): 

726 """Read an object from a pickle file specified by ButlerLocation. 

727 

728 Parameters 

729 ---------- 

730 butlerLocation : ButlerLocation 

731 The location for the object(s) to be read. 

732 

733 Returns 

734 ------- 

735 A list of objects as described by the butler location. One item for 

736 each location in butlerLocation.getLocations() 

737 """ 

738 # Create a list of Storages for the item. 

739 results = [] 

740 additionalData = butlerLocation.getAdditionalData() 

741 for locationString in butlerLocation.getLocations(): 

742 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString) 

743 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

744 if not os.path.exists(logLoc.locString()): 

745 raise RuntimeError("No such pickle file: " + logLoc.locString()) 

746 with open(logLoc.locString(), "rb") as infile: 

747 # py3: We have to specify encoding since some files were written 

748 # by python2, and 'latin1' manages that conversion safely. See: 

749 # http://stackoverflow.com/questions/28218466/unpickling-a-python-2-object-with-python-3/28218598#28218598 

750 if sys.version_info.major >= 3: 

751 finalItem = pickle.load(infile, encoding="latin1") 

752 else: 

753 finalItem = pickle.load(infile) 

754 results.append(finalItem) 

755 return results 

756 

757 

758def writePickleStorage(butlerLocation, obj): 

759 """Writes an object to a pickle file specified by ButlerLocation. 

760 

761 Parameters 

762 ---------- 

763 butlerLocation : ButlerLocation 

764 The location for the object to be written. 

765 obj : object instance 

766 The object to be written. 

767 """ 

768 additionalData = butlerLocation.getAdditionalData() 

769 locations = butlerLocation.getLocations() 

770 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString: 

771 logLoc = LogicalLocation(locationString, additionalData) 

772 with open(logLoc.locString(), "wb") as outfile: 

773 pickle.dump(obj, outfile, pickle.HIGHEST_PROTOCOL) 

774 

775 

776def readFitsCatalogStorage(butlerLocation): 

777 """Read a catalog from a FITS table specified by ButlerLocation. 

778 

779 Parameters 

780 ---------- 

781 butlerLocation : ButlerLocation 

782 The location for the object(s) to be read. 

783 

784 Returns 

785 ------- 

786 A list of objects as described by the butler location. One item for 

787 each location in butlerLocation.getLocations() 

788 """ 

789 pythonType = butlerLocation.getPythonType() 

790 if pythonType is not None: 

791 if isinstance(pythonType, str): 

792 pythonType = doImport(pythonType) 

793 results = [] 

794 additionalData = butlerLocation.getAdditionalData() 

795 for locationString in butlerLocation.getLocations(): 

796 locStringWithRoot = os.path.join(butlerLocation.getStorage().root, locationString) 

797 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

798 if not os.path.exists(logLoc.locString()): 

799 raise RuntimeError("No such FITS catalog file: " + logLoc.locString()) 

800 kwds = {} 

801 if additionalData.exists("hdu"): 

802 kwds["hdu"] = additionalData.getInt("hdu") 

803 if additionalData.exists("flags"): 

804 kwds["flags"] = additionalData.getInt("flags") 

805 finalItem = pythonType.readFits(logLoc.locString(), **kwds) 

806 results.append(finalItem) 

807 return results 

808 

809 

810def writeFitsCatalogStorage(butlerLocation, obj): 

811 """Writes a catalog to a FITS table specified by ButlerLocation. 

812 

813 Parameters 

814 ---------- 

815 butlerLocation : ButlerLocation 

816 The location for the object to be written. 

817 obj : object instance 

818 The object to be written. 

819 """ 

820 additionalData = butlerLocation.getAdditionalData() 

821 locations = butlerLocation.getLocations() 

822 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString: 

823 logLoc = LogicalLocation(locationString, additionalData) 

824 if additionalData.exists("flags"): 

825 kwds = dict(flags=additionalData.getInt("flags")) 

826 else: 

827 kwds = {} 

828 obj.writeFits(logLoc.locString(), **kwds) 

829 

830 

831def readMatplotlibStorage(butlerLocation): 

832 """Read from a butlerLocation (always fails for this storage type). 

833 

834 Parameters 

835 ---------- 

836 butlerLocation : ButlerLocation 

837 The location for the object(s) to be read. 

838 

839 Returns 

840 ------- 

841 A list of objects as described by the butler location. One item for 

842 each location in butlerLocation.getLocations() 

843 """ 

844 raise NotImplementedError("Figures saved with MatplotlibStorage cannot be retreived using the Butler.") 

845 

846 

847def writeMatplotlibStorage(butlerLocation, obj): 

848 """Writes a matplotlib.figure.Figure to a location, using the template's 

849 filename suffix to infer the file format. 

850 

851 Parameters 

852 ---------- 

853 butlerLocation : ButlerLocation 

854 The location for the object to be written. 

855 obj : matplotlib.figure.Figure 

856 The object to be written. 

857 """ 

858 additionalData = butlerLocation.getAdditionalData() 

859 locations = butlerLocation.getLocations() 

860 with SafeFilename(os.path.join(butlerLocation.getStorage().root, locations[0])) as locationString: 

861 logLoc = LogicalLocation(locationString, additionalData) 

862 # SafeFilename appends a random suffix, which corrupts the extension 

863 # matplotlib uses to guess the file format. 

864 # Instead, we extract the extension from the original location 

865 # and pass that as the format directly. 

866 _, ext = os.path.splitext(locations[0]) 

867 if ext: 

868 ext = ext[1:] # strip off leading '.' 

869 else: 

870 # If there is no extension, we let matplotlib fall back to its 

871 # default. 

872 ext = None 

873 obj.savefig(logLoc.locString(), format=ext) 

874 

875 

876def readYamlStorage(butlerLocation): 

877 """Read an object from a YAML file specified by a butlerLocation. 

878 

879 Parameters 

880 ---------- 

881 butlerLocation : ButlerLocation 

882 The location for the object(s) to be read. 

883 

884 Returns 

885 ------- 

886 A list of objects as described by the butler location. One item for 

887 each location in butlerLocation.getLocations() 

888 """ 

889 results = [] 

890 for locationString in butlerLocation.getLocations(): 

891 logLoc = LogicalLocation(butlerLocation.getStorage().locationWithRoot(locationString), 

892 butlerLocation.getAdditionalData()) 

893 if not os.path.exists(logLoc.locString()): 

894 raise RuntimeError("No such YAML file: " + logLoc.locString()) 

895 # Butler Gen2 repository configurations are handled specially 

896 if butlerLocation.pythonType == 'lsst.daf.persistence.RepositoryCfg': 

897 finalItem = Policy(filePath=logLoc.locString()) 

898 else: 

899 try: 

900 # PyYAML >=5.1 prefers a different loader 

901 loader = yaml.UnsafeLoader 

902 except AttributeError: 

903 loader = yaml.Loader 

904 with open(logLoc.locString(), "rb") as infile: 

905 finalItem = yaml.load(infile, Loader=loader) 

906 results.append(finalItem) 

907 return results 

908 

909 

910PosixStorage.registerFormatters("FitsStorage", readFitsStorage, writeFitsStorage) 

911PosixStorage.registerFormatters("ParquetStorage", readParquetStorage, writeParquetStorage) 

912PosixStorage.registerFormatters("ConfigStorage", readConfigStorage, writeConfigStorage) 

913PosixStorage.registerFormatters("PickleStorage", readPickleStorage, writePickleStorage) 

914PosixStorage.registerFormatters("FitsCatalogStorage", readFitsCatalogStorage, writeFitsCatalogStorage) 

915PosixStorage.registerFormatters("MatplotlibStorage", readMatplotlibStorage, writeMatplotlibStorage) 

916PosixStorage.registerFormatters("YamlStorage", readYamlStorage, writeYamlStorage) 

917 

918Storage.registerStorageClass(scheme='', cls=PosixStorage) 

919Storage.registerStorageClass(scheme='file', cls=PosixStorage)