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

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 pythonType = butlerLocation.getPythonType() 

532 if pythonType is not None: 

533 if isinstance(pythonType, str): 

534 pythonType = doImport(pythonType) 

535 finalItem = pythonType() 

536 finalItem.load(logLoc.locString()) 

537 results.append(finalItem) 

538 return results 

539 

540 

541def writeConfigStorage(butlerLocation, obj): 

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

543 ButlerLocation. 

544 

545 Parameters 

546 ---------- 

547 butlerLocation : ButlerLocation 

548 The location for the object to be written. 

549 obj : object instance 

550 The object to be written. 

551 """ 

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

553 with SafeFilename(filename) as locationString: 

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

555 obj.save(logLoc.locString()) 

556 

557 

558def readFitsStorage(butlerLocation): 

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

560 

561 The object is read using class or static method 

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

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

564 ``butlerLocation.getAdditionalData()``. 

565 

566 Parameters 

567 ---------- 

568 butlerLocation : ButlerLocation 

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

570 

571 Returns 

572 ------- 

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

574 each location in butlerLocation.getLocations() 

575 """ 

576 pythonType = butlerLocation.getPythonType() 

577 if pythonType is not None: 

578 if isinstance(pythonType, str): 

579 pythonType = doImport(pythonType) 

580 supportsOptions = hasattr(pythonType, "readFitsWithOptions") 

581 if not supportsOptions: 

582 from lsst.daf.base import PropertySet, PropertyList 

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

584 from lsst.afw.image import readMetadata 

585 reader = readMetadata 

586 else: 

587 reader = pythonType.readFits 

588 results = [] 

589 additionalData = butlerLocation.getAdditionalData() 

590 for locationString in butlerLocation.getLocations(): 

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

592 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

594 # because that can specify the HDU or other information 

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

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

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

598 if supportsOptions: 

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

600 else: 

601 fileName = logLoc.locString() 

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

603 

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

605 fileName = mat.group(1) 

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

607 

608 finalItem = reader(fileName, hdu=hdu) 

609 else: 

610 finalItem = reader(fileName) 

611 results.append(finalItem) 

612 return results 

613 

614 

615def writeFitsStorage(butlerLocation, obj): 

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

617 

618 The object is written using method 

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

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

621 ``butlerLocation.getAdditionalData()``. 

622 

623 Parameters 

624 ---------- 

625 butlerLocation : ButlerLocation 

626 The location for the object to be written. 

627 obj : object instance 

628 The object to be written. 

629 """ 

630 supportsOptions = hasattr(obj, "writeFitsWithOptions") 

631 additionalData = butlerLocation.getAdditionalData() 

632 locations = butlerLocation.getLocations() 

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

634 logLoc = LogicalLocation(locationString, additionalData) 

635 if supportsOptions: 

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

637 else: 

638 obj.writeFits(logLoc.locString()) 

639 

640 

641def readParquetStorage(butlerLocation): 

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

643 

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

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

646 that allows for lazy loading of the data. 

647 

648 Parameters 

649 ---------- 

650 butlerLocation : ButlerLocation 

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

652 

653 Returns 

654 ------- 

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

656 each location in butlerLocation.getLocations() 

657 """ 

658 results = [] 

659 additionalData = butlerLocation.getAdditionalData() 

660 

661 for locationString in butlerLocation.getLocations(): 

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

663 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

666 

667 pythonType = butlerLocation.getPythonType() 

668 if pythonType is not None: 

669 if isinstance(pythonType, str): 

670 pythonType = doImport(pythonType) 

671 

672 filename = logLoc.locString() 

673 

674 # pythonType will be ParquetTable (or perhaps MultilevelParquetTable) 

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

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

677 

678 return results 

679 

680 

681def writeParquetStorage(butlerLocation, obj): 

682 """Writes pandas dataframe to parquet file. 

683 

684 Parameters 

685 ---------- 

686 butlerLocation : ButlerLocation 

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

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

689 Wrapped DataFrame to write. 

690 

691 """ 

692 additionalData = butlerLocation.getAdditionalData() 

693 locations = butlerLocation.getLocations() 

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

695 logLoc = LogicalLocation(locationString, additionalData) 

696 filename = logLoc.locString() 

697 obj.write(filename) 

698 

699 

700def writeYamlStorage(butlerLocation, obj): 

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

702 

703 Parameters 

704 ---------- 

705 butlerLocation : ButlerLocation 

706 The location for the object to be written. 

707 obj : object instance 

708 The object to be written. 

709 """ 

710 additionalData = butlerLocation.getAdditionalData() 

711 locations = butlerLocation.getLocations() 

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

713 logLoc = LogicalLocation(locationString, additionalData) 

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

715 yaml.dump(obj, outfile) 

716 

717 

718def readPickleStorage(butlerLocation): 

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

720 

721 Parameters 

722 ---------- 

723 butlerLocation : ButlerLocation 

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

725 

726 Returns 

727 ------- 

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

729 each location in butlerLocation.getLocations() 

730 """ 

731 # Create a list of Storages for the item. 

732 results = [] 

733 additionalData = butlerLocation.getAdditionalData() 

734 for locationString in butlerLocation.getLocations(): 

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

736 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

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

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

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

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

743 if sys.version_info.major >= 3: 

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

745 else: 

746 finalItem = pickle.load(infile) 

747 results.append(finalItem) 

748 return results 

749 

750 

751def writePickleStorage(butlerLocation, obj): 

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

753 

754 Parameters 

755 ---------- 

756 butlerLocation : ButlerLocation 

757 The location for the object to be written. 

758 obj : object instance 

759 The object to be written. 

760 """ 

761 additionalData = butlerLocation.getAdditionalData() 

762 locations = butlerLocation.getLocations() 

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

764 logLoc = LogicalLocation(locationString, additionalData) 

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

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

767 

768 

769def readFitsCatalogStorage(butlerLocation): 

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

771 

772 Parameters 

773 ---------- 

774 butlerLocation : ButlerLocation 

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

776 

777 Returns 

778 ------- 

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

780 each location in butlerLocation.getLocations() 

781 """ 

782 pythonType = butlerLocation.getPythonType() 

783 if pythonType is not None: 

784 if isinstance(pythonType, str): 

785 pythonType = doImport(pythonType) 

786 results = [] 

787 additionalData = butlerLocation.getAdditionalData() 

788 for locationString in butlerLocation.getLocations(): 

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

790 logLoc = LogicalLocation(locStringWithRoot, additionalData) 

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

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

793 kwds = {} 

794 if additionalData.exists("hdu"): 

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

796 if additionalData.exists("flags"): 

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

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

799 results.append(finalItem) 

800 return results 

801 

802 

803def writeFitsCatalogStorage(butlerLocation, obj): 

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

805 

806 Parameters 

807 ---------- 

808 butlerLocation : ButlerLocation 

809 The location for the object to be written. 

810 obj : object instance 

811 The object to be written. 

812 """ 

813 additionalData = butlerLocation.getAdditionalData() 

814 locations = butlerLocation.getLocations() 

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

816 logLoc = LogicalLocation(locationString, additionalData) 

817 if additionalData.exists("flags"): 

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

819 else: 

820 kwds = {} 

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

822 

823 

824def readMatplotlibStorage(butlerLocation): 

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

826 

827 Parameters 

828 ---------- 

829 butlerLocation : ButlerLocation 

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

831 

832 Returns 

833 ------- 

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

835 each location in butlerLocation.getLocations() 

836 """ 

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

838 

839 

840def writeMatplotlibStorage(butlerLocation, obj): 

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

842 filename suffix to infer the file format. 

843 

844 Parameters 

845 ---------- 

846 butlerLocation : ButlerLocation 

847 The location for the object to be written. 

848 obj : matplotlib.figure.Figure 

849 The object to be written. 

850 """ 

851 additionalData = butlerLocation.getAdditionalData() 

852 locations = butlerLocation.getLocations() 

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

854 logLoc = LogicalLocation(locationString, additionalData) 

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

856 # matplotlib uses to guess the file format. 

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

858 # and pass that as the format directly. 

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

860 if ext: 

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

862 else: 

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

864 # default. 

865 ext = None 

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

867 

868 

869def readPafStorage(butlerLocation): 

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

871 

872 Parameters 

873 ---------- 

874 butlerLocation : ButlerLocation 

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

876 

877 Returns 

878 ------- 

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

880 each location in butlerLocation.getLocations() 

881 """ 

882 results = [] 

883 for locationString in butlerLocation.getLocations(): 

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

885 butlerLocation.getAdditionalData()) 

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

887 results.append(finalItem) 

888 return results 

889 

890 

891def readYamlStorage(butlerLocation): 

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

893 

894 Parameters 

895 ---------- 

896 butlerLocation : ButlerLocation 

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

898 

899 Returns 

900 ------- 

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

902 each location in butlerLocation.getLocations() 

903 """ 

904 results = [] 

905 for locationString in butlerLocation.getLocations(): 

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

907 butlerLocation.getAdditionalData()) 

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

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

910 # Butler Gen2 repository configurations are handled specially 

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

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

913 else: 

914 try: 

915 # PyYAML >=5.1 prefers a different loader 

916 loader = yaml.FullLoader 

917 except AttributeError: 

918 loader = yaml.Loader 

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

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

921 results.append(finalItem) 

922 return results 

923 

924 

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

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

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

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

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

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

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

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

933 

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

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