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 

38import lsst.pex.policy as pexPolicy 

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', '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 

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

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

348 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 path.startswith(rootDir + "/"): 

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

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

451 pathPrefix = rootDir 

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

453 path = path[1:] 

454 pathPrefix = None 

455 else: 

456 # Search for prefix that is the same as root 

457 pathPrefix = os.path.dirname(path) 

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

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

460 break 

461 pathPrefix = os.path.dirname(pathPrefix) 

462 if pathPrefix == "/": 

463 path = path[1:] 

464 elif pathPrefix != "": 

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

466 

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

468 # Strip off any cfitsio bracketed extension if present 

469 strippedPath = path 

470 pathStripped = None 

471 firstBracket = path.find("[") 

472 if firstBracket != -1: 

473 strippedPath = path[:firstBracket] 

474 pathStripped = path[firstBracket:] 

475 

476 dir = rootDir 

477 while True: 

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

479 if len(paths) > 0: 

480 if pathPrefix != rootDir: 

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

482 if pathStripped is not None: 

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

484 return paths 

485 if searchParents: 

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

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

488 return None 

489 else: 

490 return None 

491 

492 @staticmethod 

493 def storageExists(uri): 

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

495 

496 Parameters 

497 ---------- 

498 root : string 

499 URI to the the root location of the storage 

500 

501 Returns 

502 ------- 

503 bool 

504 True if the storage exists, false if not 

505 """ 

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

507 

508 

509def readConfigStorage(butlerLocation): 

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

511 

512 Parameters 

513 ---------- 

514 butlerLocation : ButlerLocation 

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

516 

517 Returns 

518 ------- 

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

520 each location in butlerLocation.getLocations() 

521 """ 

522 results = [] 

523 for locationString in butlerLocation.getLocations(): 

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

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

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

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

528 pythonType = butlerLocation.getPythonType() 

529 if pythonType is not None: 

530 if isinstance(pythonType, str): 

531 pythonType = doImport(pythonType) 

532 finalItem = pythonType() 

533 finalItem.load(logLoc.locString()) 

534 results.append(finalItem) 

535 return results 

536 

537 

538def writeConfigStorage(butlerLocation, obj): 

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

540 ButlerLocation. 

541 

542 Parameters 

543 ---------- 

544 butlerLocation : ButlerLocation 

545 The location for the object to be written. 

546 obj : object instance 

547 The object to be written. 

548 """ 

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

550 with SafeFilename(filename) as locationString: 

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

552 obj.save(logLoc.locString()) 

553 

554 

555def readFitsStorage(butlerLocation): 

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

557 

558 The object is read using class or static method 

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

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

561 ``butlerLocation.getAdditionalData()``. 

562 

563 Parameters 

564 ---------- 

565 butlerLocation : ButlerLocation 

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

567 

568 Returns 

569 ------- 

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

571 each location in butlerLocation.getLocations() 

572 """ 

573 pythonType = butlerLocation.getPythonType() 

574 if pythonType is not None: 

575 if isinstance(pythonType, str): 

576 pythonType = doImport(pythonType) 

577 supportsOptions = hasattr(pythonType, "readFitsWithOptions") 

578 if not supportsOptions: 

579 from lsst.daf.base import PropertySet, PropertyList 

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

581 from lsst.afw.image import readMetadata 

582 reader = readMetadata 

583 else: 

584 reader = pythonType.readFits 

585 results = [] 

586 additionalData = butlerLocation.getAdditionalData() 

587 for locationString in butlerLocation.getLocations(): 

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

589 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

591 # because that can specify the HDU or other information 

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

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

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

595 if supportsOptions: 

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

597 else: 

598 fileName = logLoc.locString() 

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

600 

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

602 fileName = mat.group(1) 

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

604 

605 finalItem = reader(fileName, hdu=hdu) 

606 else: 

607 finalItem = reader(fileName) 

608 results.append(finalItem) 

609 return results 

610 

611 

612def writeFitsStorage(butlerLocation, obj): 

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

614 

615 The object is written using method 

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

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

618 ``butlerLocation.getAdditionalData()``. 

619 

620 Parameters 

621 ---------- 

622 butlerLocation : ButlerLocation 

623 The location for the object to be written. 

624 obj : object instance 

625 The object to be written. 

626 """ 

627 supportsOptions = hasattr(obj, "writeFitsWithOptions") 

628 additionalData = butlerLocation.getAdditionalData() 

629 locations = butlerLocation.getLocations() 

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

631 logLoc = LogicalLocation(locationString, additionalData) 

632 if supportsOptions: 

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

634 else: 

635 obj.writeFits(logLoc.locString()) 

636 

637 

638def readParquetStorage(butlerLocation): 

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

640 

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

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

643 that allows for lazy loading of the data. 

644 

645 Parameters 

646 ---------- 

647 butlerLocation : ButlerLocation 

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

649 

650 Returns 

651 ------- 

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

653 each location in butlerLocation.getLocations() 

654 """ 

655 results = [] 

656 additionalData = butlerLocation.getAdditionalData() 

657 

658 for locationString in butlerLocation.getLocations(): 

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

660 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

663 

664 pythonType = butlerLocation.getPythonType() 

665 if pythonType is not None: 

666 if isinstance(pythonType, str): 

667 pythonType = doImport(pythonType) 

668 

669 filename = logLoc.locString() 

670 

671 # pythonType will be ParquetTable (or perhaps MultilevelParquetTable) 

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

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

674 

675 return results 

676 

677 

678def writeParquetStorage(butlerLocation, obj): 

679 """Writes pandas dataframe to parquet file. 

680 

681 Parameters 

682 ---------- 

683 butlerLocation : ButlerLocation 

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

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

686 Wrapped DataFrame to write. 

687 

688 """ 

689 additionalData = butlerLocation.getAdditionalData() 

690 locations = butlerLocation.getLocations() 

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

692 logLoc = LogicalLocation(locationString, additionalData) 

693 filename = logLoc.locString() 

694 obj.write(filename) 

695 

696 

697def writeYamlStorage(butlerLocation, obj): 

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

699 

700 Parameters 

701 ---------- 

702 butlerLocation : ButlerLocation 

703 The location for the object to be written. 

704 obj : object instance 

705 The object to be written. 

706 """ 

707 additionalData = butlerLocation.getAdditionalData() 

708 locations = butlerLocation.getLocations() 

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

710 logLoc = LogicalLocation(locationString, additionalData) 

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

712 yaml.dump(obj, outfile) 

713 

714 

715def readPickleStorage(butlerLocation): 

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

717 

718 Parameters 

719 ---------- 

720 butlerLocation : ButlerLocation 

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

722 

723 Returns 

724 ------- 

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

726 each location in butlerLocation.getLocations() 

727 """ 

728 # Create a list of Storages for the item. 

729 results = [] 

730 additionalData = butlerLocation.getAdditionalData() 

731 for locationString in butlerLocation.getLocations(): 

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

733 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

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

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

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

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

740 if sys.version_info.major >= 3: 

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

742 else: 

743 finalItem = pickle.load(infile) 

744 results.append(finalItem) 

745 return results 

746 

747 

748def writePickleStorage(butlerLocation, obj): 

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

750 

751 Parameters 

752 ---------- 

753 butlerLocation : ButlerLocation 

754 The location for the object to be written. 

755 obj : object instance 

756 The object to be written. 

757 """ 

758 additionalData = butlerLocation.getAdditionalData() 

759 locations = butlerLocation.getLocations() 

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

761 logLoc = LogicalLocation(locationString, additionalData) 

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

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

764 

765 

766def readFitsCatalogStorage(butlerLocation): 

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

768 

769 Parameters 

770 ---------- 

771 butlerLocation : ButlerLocation 

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

773 

774 Returns 

775 ------- 

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

777 each location in butlerLocation.getLocations() 

778 """ 

779 pythonType = butlerLocation.getPythonType() 

780 if pythonType is not None: 

781 if isinstance(pythonType, str): 

782 pythonType = doImport(pythonType) 

783 results = [] 

784 additionalData = butlerLocation.getAdditionalData() 

785 for locationString in butlerLocation.getLocations(): 

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

787 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

790 kwds = {} 

791 if additionalData.exists("hdu"): 

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

793 if additionalData.exists("flags"): 

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

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

796 results.append(finalItem) 

797 return results 

798 

799 

800def writeFitsCatalogStorage(butlerLocation, obj): 

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

802 

803 Parameters 

804 ---------- 

805 butlerLocation : ButlerLocation 

806 The location for the object to be written. 

807 obj : object instance 

808 The object to be written. 

809 """ 

810 additionalData = butlerLocation.getAdditionalData() 

811 locations = butlerLocation.getLocations() 

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

813 logLoc = LogicalLocation(locationString, additionalData) 

814 if additionalData.exists("flags"): 

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

816 else: 

817 kwds = {} 

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

819 

820 

821def readMatplotlibStorage(butlerLocation): 

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

823 

824 Parameters 

825 ---------- 

826 butlerLocation : ButlerLocation 

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

828 

829 Returns 

830 ------- 

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

832 each location in butlerLocation.getLocations() 

833 """ 

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

835 

836 

837def writeMatplotlibStorage(butlerLocation, obj): 

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

839 filename suffix to infer the file format. 

840 

841 Parameters 

842 ---------- 

843 butlerLocation : ButlerLocation 

844 The location for the object to be written. 

845 obj : matplotlib.figure.Figure 

846 The object to be written. 

847 """ 

848 additionalData = butlerLocation.getAdditionalData() 

849 locations = butlerLocation.getLocations() 

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

851 logLoc = LogicalLocation(locationString, additionalData) 

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

853 # matplotlib uses to guess the file format. 

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

855 # and pass that as the format directly. 

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

857 if ext: 

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

859 else: 

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

861 # default. 

862 ext = None 

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

864 

865 

866def readPafStorage(butlerLocation): 

867 """Read a policy from a PAF file specified by a ButlerLocation. 

868 

869 Parameters 

870 ---------- 

871 butlerLocation : ButlerLocation 

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

873 

874 Returns 

875 ------- 

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

877 each location in butlerLocation.getLocations() 

878 """ 

879 results = [] 

880 for locationString in butlerLocation.getLocations(): 

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

882 butlerLocation.getAdditionalData()) 

883 finalItem = pexPolicy.Policy.createPolicy(logLoc.locString()) 

884 results.append(finalItem) 

885 return results 

886 

887 

888def readYamlStorage(butlerLocation): 

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

890 

891 Parameters 

892 ---------- 

893 butlerLocation : ButlerLocation 

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

895 

896 Returns 

897 ------- 

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

899 each location in butlerLocation.getLocations() 

900 """ 

901 results = [] 

902 for locationString in butlerLocation.getLocations(): 

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

904 butlerLocation.getAdditionalData()) 

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

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

907 # Butler Gen2 repository configurations are handled specially 

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

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

910 else: 

911 try: 

912 # PyYAML >=5.1 prefers a different loader 

913 loader = yaml.FullLoader 

914 except AttributeError: 

915 loader = yaml.Loader 

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

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

918 results.append(finalItem) 

919 return results 

920 

921 

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

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

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

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

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

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

928PosixStorage.registerFormatters("PafStorage", readFormatter=readPafStorage) 

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

930 

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

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