Hide keyboard shortcuts

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 

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 

38from .safeFileIo import SafeFilename, safeMakeDir 

39 

40 

41__all__ = ["PosixStorage"] 

42 

43 

44class PosixStorage(StorageInterface): 

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

46 

47 Parameters 

48 ---------- 

49 uri : string 

50 URI or path that is used as the storage location. 

51 create : bool 

52 If True a new repository will be created at the root location if it 

53 does not exist. If False then a new repository will not be created. 

54 

55 Raises 

56 ------ 

57 NoRepositroyAtRoot 

58 If create is False and a repository does not exist at the root 

59 specified by uri then NoRepositroyAtRoot is raised. 

60 """ 

61 

62 def __init__(self, uri, create): 

63 self.log = Log.getLogger("daf.persistence.butler") 

64 self.root = self._pathFromURI(uri) 

65 if self.root and not os.path.exists(self.root): 

66 if not create: 

67 raise NoRepositroyAtRoot("No repository at {}".format(uri)) 

68 safeMakeDir(self.root) 

69 

70 def __repr__(self): 

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

72 

73 @staticmethod 

74 def _pathFromURI(uri): 

75 """Get the path part of the URI""" 

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

77 

78 @staticmethod 

79 def relativePath(fromPath, toPath): 

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

81 

82 Parameters 

83 ---------- 

84 fromPath : string 

85 A path at which to start. It can be a relative path or an 

86 absolute path. 

87 toPath : string 

88 A target location. It can be a relative path or an absolute path. 

89 

90 Returns 

91 ------- 

92 string 

93 A relative path that describes the path from fromPath to toPath. 

94 """ 

95 fromPath = os.path.realpath(fromPath) 

96 return os.path.relpath(toPath, fromPath) 

97 

98 @staticmethod 

99 def absolutePath(fromPath, relativePath): 

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

101 

102 Parameters 

103 ---------- 

104 fromPath : the starting location 

105 A location at which to start. It can be a relative path or an 

106 absolute path. 

107 relativePath : the location relative to fromPath 

108 A relative path. 

109 

110 Returns 

111 ------- 

112 string 

113 Path that is an absolute path representation of fromPath + 

114 relativePath, if one exists. If relativePath is absolute or if 

115 fromPath is not related to relativePath then relativePath will be 

116 returned. 

117 """ 

118 if os.path.isabs(relativePath): 

119 return relativePath 

120 fromPath = os.path.realpath(fromPath) 

121 return os.path.normpath(os.path.join(fromPath, relativePath)) 

122 

123 @staticmethod 

124 def getRepositoryCfg(uri): 

125 """Get a persisted RepositoryCfg 

126 

127 Parameters 

128 ---------- 

129 uri : URI or path to a RepositoryCfg 

130 Description 

131 

132 Returns 

133 ------- 

134 A RepositoryCfg instance or None 

135 """ 

136 storage = Storage.makeFromURI(uri) 

137 location = ButlerLocation(pythonType=RepositoryCfg, 

138 cppType=None, 

139 storageName=None, 

140 locationList='repositoryCfg.yaml', 

141 dataId={}, 

142 mapper=None, 

143 storage=storage, 

144 usedDataId=None, 

145 datasetType=None) 

146 return storage.read(location) 

147 

148 @staticmethod 

149 def putRepositoryCfg(cfg, loc=None): 

150 storage = Storage.makeFromURI(cfg.root if loc is None else loc, create=True) 

151 location = ButlerLocation(pythonType=RepositoryCfg, 

152 cppType=None, 

153 storageName=None, 

154 locationList='repositoryCfg.yaml', 

155 dataId={}, 

156 mapper=None, 

157 storage=storage, 

158 usedDataId=None, 

159 datasetType=None) 

160 storage.write(location, cfg) 

161 

162 @staticmethod 

163 def getMapperClass(root): 

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

165 

166 Supports the legacy _parent symlink search (which was only ever posix-only. This should not be used by 

167 new code and repositories; they should use the Repository parentCfg mechanism. 

168 

169 Parameters 

170 ---------- 

171 root : string 

172 The location of a persisted ReositoryCfg is (new style repos), or 

173 the location where a _mapper file is (old style repos). 

174 

175 Returns 

176 ------- 

177 A class object or a class instance, depending on the state of the 

178 mapper when the repository was created. 

179 """ 

180 if not (root): 

181 return None 

182 

183 cfg = PosixStorage.getRepositoryCfg(root) 

184 if cfg is not None: 

185 return cfg.mapper 

186 

187 # Find a "_mapper" file containing the mapper class name 

188 basePath = root 

189 mapperFile = "_mapper" 

190 while not os.path.exists(os.path.join(basePath, mapperFile)): 

191 # Break abstraction by following _parent links from CameraMapper 

192 if os.path.exists(os.path.join(basePath, "_parent")): 

193 basePath = os.path.join(basePath, "_parent") 

194 else: 

195 mapperFile = None 

196 break 

197 

198 if mapperFile is not None: 

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

200 

201 # Read the name of the mapper class and instantiate it 

202 with open(mapperFile, "r") as f: 

203 mapperName = f.readline().strip() 

204 components = mapperName.split(".") 

205 if len(components) <= 1: 

206 raise RuntimeError("Unqualified mapper name %s in %s" % 

207 (mapperName, mapperFile)) 

208 pkg = importlib.import_module(".".join(components[:-1])) 

209 return getattr(pkg, components[-1]) 

210 

211 return None 

212 

213 @staticmethod 

214 def getParentSymlinkPath(root): 

215 """For Butler V1 Repositories only, if a _parent symlink exists, get the location pointed to by the 

216 symlink. 

217 

218 Parameters 

219 ---------- 

220 root : string 

221 A path to the folder on the local filesystem. 

222 

223 Returns 

224 ------- 

225 string or None 

226 A path to the parent folder indicated by the _parent symlink, or None if there is no _parent 

227 symlink at root. 

228 """ 

229 linkpath = os.path.join(root, '_parent') 

230 if os.path.exists(linkpath): 

231 try: 

232 return os.readlink(os.path.join(root, '_parent')) 

233 except OSError: 

234 # some of the unit tests rely on a folder called _parent instead of a symlink to aother 

235 # location. Allow that; return the path of that folder. 

236 return os.path.join(root, '_parent') 

237 return None 

238 

239 def write(self, butlerLocation, obj): 

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

241 ButlerLocation 

242 

243 Parameters 

244 ---------- 

245 butlerLocation : ButlerLocation 

246 The location & formatting for the object to be written. 

247 obj : object instance 

248 The object to be written. 

249 """ 

250 self.log.debug("Put location=%s obj=%s", butlerLocation, obj) 

251 

252 writeFormatter = self.getWriteFormatter(butlerLocation.getStorageName()) 

253 if not writeFormatter: 

254 writeFormatter = self.getWriteFormatter(butlerLocation.getPythonType()) 

255 if writeFormatter: 

256 writeFormatter(butlerLocation, obj) 

257 return 

258 

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

260 

261 def read(self, butlerLocation): 

262 """Read from a butlerLocation. 

263 

264 Parameters 

265 ---------- 

266 butlerLocation : ButlerLocation 

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

268 

269 Returns 

270 ------- 

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

272 each location in butlerLocation.getLocations() 

273 """ 

274 readFormatter = self.getReadFormatter(butlerLocation.getStorageName()) 

275 if not readFormatter: 

276 readFormatter = self.getReadFormatter(butlerLocation.getPythonType()) 

277 if readFormatter: 

278 return readFormatter(butlerLocation) 

279 

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

281 

282 def butlerLocationExists(self, location): 

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

284 """ 

285 storageName = location.getStorageName() 

286 if storageName not in ('FitsStorage', 

287 'PickleStorage', 'ConfigStorage', 'FitsCatalogStorage', 

288 'YamlStorage', 'ParquetStorage', 'MatplotlibStorage'): 

289 self.log.warn("butlerLocationExists for non-supported storage %s" % location) 

290 return False 

291 for locationString in location.getLocations(): 

292 logLoc = LogicalLocation(locationString, location.getAdditionalData()).locString() 

293 obj = self.instanceSearch(path=logLoc) 

294 if obj: 

295 return True 

296 return False 

297 

298 def exists(self, location): 

299 """Check if location exists. 

300 

301 Parameters 

302 ---------- 

303 location : ButlerLocation or string 

304 A a string or a ButlerLocation that describes the location of an 

305 object in this storage. 

306 

307 Returns 

308 ------- 

309 bool 

310 True if exists, else False. 

311 """ 

312 if isinstance(location, ButlerLocation): 

313 return self.butlerLocationExists(location) 

314 

315 obj = self.instanceSearch(path=location) 

316 return bool(obj) 

317 

318 def locationWithRoot(self, location): 

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

320 

321 :param location: 

322 :return: 

323 """ 

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

325 

326 @staticmethod 

327 def v1RepoExists(root): 

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

329 

330 Version 1 Repositories only exist in posix storages, do not have a 

331 RepositoryCfg file, and contain either a registry.sqlite3 file, a 

332 _mapper file, or a _parent link. 

333 

334 Parameters 

335 ---------- 

336 root : string 

337 A path to a folder on the local filesystem. 

338 

339 Returns 

340 ------- 

341 bool 

342 True if the repository at root exists, else False. 

343 """ 

344 return os.path.exists(root) and ( 

345 os.path.exists(os.path.join(root, "registry.sqlite3")) 

346 or os.path.exists(os.path.join(root, "_mapper")) 

347 or os.path.exists(os.path.join(root, "_parent")) 

348 ) 

349 

350 def copyFile(self, fromLocation, toLocation): 

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

352 

353 Parameters 

354 ---------- 

355 fromLocation : path 

356 Path and name of existing file. 

357 toLocation : path 

358 Path and name of new file. 

359 

360 Returns 

361 ------- 

362 None 

363 """ 

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

365 

366 def getLocalFile(self, path): 

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

368 temporary if needed. 

369 

370 Parameters 

371 ---------- 

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

373 

374 Returns 

375 ------- 

376 A handle to a local copy of the file. If storage is remote it will be 

377 a temporary file. If storage is local it may be the original file or 

378 a temporary file. The file name can be gotten via the 'name' property 

379 of the returned object. 

380 """ 

381 p = os.path.join(self.root, path) 

382 try: 

383 return open(p) 

384 except IOError as e: 

385 if e.errno == 2: # 'No such file or directory' 

386 return None 

387 else: 

388 raise e 

389 

390 def instanceSearch(self, path): 

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

392 

393 If the path contains an HDU indicator (a number in brackets before the 

394 dot, e.g. 'foo.fits[1]', this will be stripped when searching and so 

395 will match filenames without the HDU indicator, e.g. 'foo.fits'. The 

396 path returned WILL contain the indicator though, e.g. ['foo.fits[1]']. 

397 

398 Parameters 

399 ---------- 

400 path : string 

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

402 

403 Returns 

404 ------- 

405 string or None 

406 The location that was found, or None if no location was found. 

407 """ 

408 return self.search(self.root, path) 

409 

410 @staticmethod 

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

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

413 

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

415 following the Butler v1 _parent symlink 

416 

417 If the path contains an HDU indicator (a number in brackets, e.g. 

418 'foo.fits[1]', this will be stripped when searching and so 

419 will match filenames without the HDU indicator, e.g. 'foo.fits'. The 

420 path returned WILL contain the indicator though, e.g. ['foo.fits[1]']. 

421 

422 Parameters 

423 ---------- 

424 root : string 

425 The path to the root directory. 

426 path : string 

427 The path to the file within the root directory. 

428 searchParents : bool, optional 

429 For Butler v1 repositories only, if true and a _parent symlink 

430 exists, then the directory at _parent will be searched if the file 

431 is not found in the root repository. Will continue searching the 

432 parent of the parent until the file is found or no additional 

433 parent exists. 

434 

435 Returns 

436 ------- 

437 string or None 

438 The location that was found, or None if no location was found. 

439 """ 

440 # Separate path into a root-equivalent prefix (in dir) and the rest 

441 # (left in path) 

442 rootDir = root 

443 # First remove trailing slashes (#2527) 

444 while len(rootDir) > 1 and rootDir[-1] == '/': 

445 rootDir = rootDir[:-1] 

446 

447 if not path.startswith('/'): 

448 # Most common case is a relative path from a template 

449 pathPrefix = None 

450 elif path.startswith(rootDir + "/"): 

451 # Common case; we have the same root prefix string 

452 path = path[len(rootDir + '/'):] 

453 pathPrefix = rootDir 

454 elif rootDir == "/" and path.startswith("/"): 

455 path = path[1:] 

456 pathPrefix = None 

457 else: 

458 # Search for prefix that is the same as root 

459 pathPrefix = os.path.dirname(path) 

460 while pathPrefix != "" and pathPrefix != "/": 

461 if os.path.realpath(pathPrefix) == os.path.realpath(root): 

462 break 

463 pathPrefix = os.path.dirname(pathPrefix) 

464 if pathPrefix == "/": 

465 path = path[1:] 

466 elif pathPrefix != "": 

467 path = path[len(pathPrefix)+1:] 

468 

469 # Now search for the path in the root or its parents 

470 # Strip off any cfitsio bracketed extension if present 

471 strippedPath = path 

472 pathStripped = None 

473 firstBracket = path.find("[") 

474 if firstBracket != -1: 

475 strippedPath = path[:firstBracket] 

476 pathStripped = path[firstBracket:] 

477 

478 dir = rootDir 

479 while True: 

480 paths = glob.glob(os.path.join(dir, strippedPath)) 

481 if len(paths) > 0: 

482 if pathPrefix != rootDir: 

483 paths = [p[len(rootDir+'/'):] for p in paths] 

484 if pathStripped is not None: 

485 paths = [p + pathStripped for p in paths] 

486 return paths 

487 if searchParents: 

488 dir = os.path.join(dir, "_parent") 

489 if not os.path.exists(dir): 

490 return None 

491 else: 

492 return None 

493 

494 @staticmethod 

495 def storageExists(uri): 

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

497 

498 Parameters 

499 ---------- 

500 root : string 

501 URI to the the root location of the storage 

502 

503 Returns 

504 ------- 

505 bool 

506 True if the storage exists, false if not 

507 """ 

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

509 

510 

511def readConfigStorage(butlerLocation): 

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

513 

514 Parameters 

515 ---------- 

516 butlerLocation : ButlerLocation 

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

518 

519 Returns 

520 ------- 

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

522 each location in butlerLocation.getLocations() 

523 """ 

524 results = [] 

525 for locationString in butlerLocation.getLocations(): 

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

527 logLoc = LogicalLocation(locStringWithRoot, butlerLocation.getAdditionalData()) 

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

529 raise RuntimeError("No such config file: " + logLoc.locString()) 

530 pythonType = butlerLocation.getPythonType() 

531 if pythonType is not None: 

532 if isinstance(pythonType, str): 

533 pythonType = doImport(pythonType) 

534 finalItem = pythonType() 

535 finalItem.load(logLoc.locString()) 

536 results.append(finalItem) 

537 return results 

538 

539 

540def writeConfigStorage(butlerLocation, obj): 

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

542 ButlerLocation. 

543 

544 Parameters 

545 ---------- 

546 butlerLocation : ButlerLocation 

547 The location for the object to be written. 

548 obj : object instance 

549 The object to be written. 

550 """ 

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

552 with SafeFilename(filename) as locationString: 

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

554 obj.save(logLoc.locString()) 

555 

556 

557def readFitsStorage(butlerLocation): 

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

559 

560 The object is read using class or static method 

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

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

563 ``butlerLocation.getAdditionalData()``. 

564 

565 Parameters 

566 ---------- 

567 butlerLocation : ButlerLocation 

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

569 

570 Returns 

571 ------- 

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

573 each location in butlerLocation.getLocations() 

574 """ 

575 pythonType = butlerLocation.getPythonType() 

576 if pythonType is not None: 

577 if isinstance(pythonType, str): 

578 pythonType = doImport(pythonType) 

579 supportsOptions = hasattr(pythonType, "readFitsWithOptions") 

580 if not supportsOptions: 

581 from lsst.daf.base import PropertySet, PropertyList 

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

583 from lsst.afw.fits import readMetadata 

584 reader = readMetadata 

585 else: 

586 reader = pythonType.readFits 

587 results = [] 

588 additionalData = butlerLocation.getAdditionalData() 

589 for locationString in butlerLocation.getLocations(): 

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

591 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

593 # because that can specify the HDU or other information 

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

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

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

597 if supportsOptions: 

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

599 else: 

600 fileName = logLoc.locString() 

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

602 

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

604 fileName = mat.group(1) 

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

606 

607 finalItem = reader(fileName, hdu=hdu) 

608 else: 

609 finalItem = reader(fileName) 

610 results.append(finalItem) 

611 return results 

612 

613 

614def writeFitsStorage(butlerLocation, obj): 

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

616 

617 The object is written using method 

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

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

620 ``butlerLocation.getAdditionalData()``. 

621 

622 Parameters 

623 ---------- 

624 butlerLocation : ButlerLocation 

625 The location for the object to be written. 

626 obj : object instance 

627 The object to be written. 

628 """ 

629 supportsOptions = hasattr(obj, "writeFitsWithOptions") 

630 additionalData = butlerLocation.getAdditionalData() 

631 locations = butlerLocation.getLocations() 

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

633 logLoc = LogicalLocation(locationString, additionalData) 

634 if supportsOptions: 

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

636 else: 

637 obj.writeFits(logLoc.locString()) 

638 

639 

640def readParquetStorage(butlerLocation): 

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

642 

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

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

645 that allows for lazy loading of the data. 

646 

647 Parameters 

648 ---------- 

649 butlerLocation : ButlerLocation 

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

651 

652 Returns 

653 ------- 

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

655 each location in butlerLocation.getLocations() 

656 """ 

657 results = [] 

658 additionalData = butlerLocation.getAdditionalData() 

659 

660 for locationString in butlerLocation.getLocations(): 

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

662 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

665 

666 pythonType = butlerLocation.getPythonType() 

667 if pythonType is not None: 

668 if isinstance(pythonType, str): 

669 pythonType = doImport(pythonType) 

670 

671 filename = logLoc.locString() 

672 

673 # pythonType will be ParquetTable (or perhaps MultilevelParquetTable) 

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

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

676 

677 return results 

678 

679 

680def writeParquetStorage(butlerLocation, obj): 

681 """Writes pandas dataframe to parquet file. 

682 

683 Parameters 

684 ---------- 

685 butlerLocation : ButlerLocation 

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

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

688 Wrapped DataFrame to write. 

689 

690 """ 

691 additionalData = butlerLocation.getAdditionalData() 

692 locations = butlerLocation.getLocations() 

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

694 logLoc = LogicalLocation(locationString, additionalData) 

695 filename = logLoc.locString() 

696 obj.write(filename) 

697 

698 

699def writeYamlStorage(butlerLocation, obj): 

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

701 

702 Parameters 

703 ---------- 

704 butlerLocation : ButlerLocation 

705 The location for the object to be written. 

706 obj : object instance 

707 The object to be written. 

708 """ 

709 additionalData = butlerLocation.getAdditionalData() 

710 locations = butlerLocation.getLocations() 

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

712 logLoc = LogicalLocation(locationString, additionalData) 

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

714 yaml.dump(obj, outfile) 

715 

716 

717def readPickleStorage(butlerLocation): 

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

719 

720 Parameters 

721 ---------- 

722 butlerLocation : ButlerLocation 

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

724 

725 Returns 

726 ------- 

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

728 each location in butlerLocation.getLocations() 

729 """ 

730 # Create a list of Storages for the item. 

731 results = [] 

732 additionalData = butlerLocation.getAdditionalData() 

733 for locationString in butlerLocation.getLocations(): 

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

735 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

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

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

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

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

742 if sys.version_info.major >= 3: 

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

744 else: 

745 finalItem = pickle.load(infile) 

746 results.append(finalItem) 

747 return results 

748 

749 

750def writePickleStorage(butlerLocation, obj): 

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

752 

753 Parameters 

754 ---------- 

755 butlerLocation : ButlerLocation 

756 The location for the object to be written. 

757 obj : object instance 

758 The object to be written. 

759 """ 

760 additionalData = butlerLocation.getAdditionalData() 

761 locations = butlerLocation.getLocations() 

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

763 logLoc = LogicalLocation(locationString, additionalData) 

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

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

766 

767 

768def readFitsCatalogStorage(butlerLocation): 

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

770 

771 Parameters 

772 ---------- 

773 butlerLocation : ButlerLocation 

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

775 

776 Returns 

777 ------- 

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

779 each location in butlerLocation.getLocations() 

780 """ 

781 pythonType = butlerLocation.getPythonType() 

782 if pythonType is not None: 

783 if isinstance(pythonType, str): 

784 pythonType = doImport(pythonType) 

785 results = [] 

786 additionalData = butlerLocation.getAdditionalData() 

787 for locationString in butlerLocation.getLocations(): 

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

789 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

792 kwds = {} 

793 if additionalData.exists("hdu"): 

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

795 if additionalData.exists("flags"): 

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

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

798 results.append(finalItem) 

799 return results 

800 

801 

802def writeFitsCatalogStorage(butlerLocation, obj): 

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

804 

805 Parameters 

806 ---------- 

807 butlerLocation : ButlerLocation 

808 The location for the object to be written. 

809 obj : object instance 

810 The object to be written. 

811 """ 

812 additionalData = butlerLocation.getAdditionalData() 

813 locations = butlerLocation.getLocations() 

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

815 logLoc = LogicalLocation(locationString, additionalData) 

816 if additionalData.exists("flags"): 

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

818 else: 

819 kwds = {} 

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

821 

822 

823def readMatplotlibStorage(butlerLocation): 

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

825 

826 Parameters 

827 ---------- 

828 butlerLocation : ButlerLocation 

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

830 

831 Returns 

832 ------- 

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

834 each location in butlerLocation.getLocations() 

835 """ 

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

837 

838 

839def writeMatplotlibStorage(butlerLocation, obj): 

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

841 filename suffix to infer the file format. 

842 

843 Parameters 

844 ---------- 

845 butlerLocation : ButlerLocation 

846 The location for the object to be written. 

847 obj : matplotlib.figure.Figure 

848 The object to be written. 

849 """ 

850 additionalData = butlerLocation.getAdditionalData() 

851 locations = butlerLocation.getLocations() 

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

853 logLoc = LogicalLocation(locationString, additionalData) 

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

855 # matplotlib uses to guess the file format. 

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

857 # and pass that as the format directly. 

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

859 if ext: 

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

861 else: 

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

863 # default. 

864 ext = None 

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

866 

867 

868def readYamlStorage(butlerLocation): 

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

870 

871 Parameters 

872 ---------- 

873 butlerLocation : ButlerLocation 

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

875 

876 Returns 

877 ------- 

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

879 each location in butlerLocation.getLocations() 

880 """ 

881 results = [] 

882 for locationString in butlerLocation.getLocations(): 

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

884 butlerLocation.getAdditionalData()) 

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

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

887 # Butler Gen2 repository configurations are handled specially 

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

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

890 else: 

891 try: 

892 # PyYAML >=5.1 prefers a different loader 

893 loader = yaml.FullLoader 

894 except AttributeError: 

895 loader = yaml.Loader 

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

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

898 results.append(finalItem) 

899 return results 

900 

901 

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

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

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

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

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

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

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

909 

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

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