Coverage for pesummary/gw/gracedb.py: 76.5%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3from ligo.gracedb.rest import GraceDb 

4from ligo.gracedb.exceptions import HTTPError 

5from pesummary.utils.utils import logger 

6from pesummary import conf 

7 

8__author__ = ["Charlie Hoy <charlie.hoy@ligo.org>"] 

9 

10 

11def _get_gracedb_client(service_url=conf.gracedb_server): 

12 """Return a ligo.gracedb.rest.GraceDb object 

13 

14 Parameters 

15 ---------- 

16 service_url: str, optional 

17 service url you wish to use when accessing data from GraceDB 

18 """ 

19 return GraceDb(service_url=service_url) 

20 

21 

22def get_gracedb_file( 

23 gracedb_id, filename, service_url=conf.gracedb_server 

24): 

25 """Grab a file from GraceDB for a specific event 

26 

27 Parameters 

28 ---------- 

29 gracedb_id: str 

30 the GraceDB id of the event you wish to retrieve the data for 

31 filename: str 

32 the name of the file you wish to retrieve from GraceDB 

33 service_url: str, optional 

34 service url you wish to use when accessing data from GraceDB 

35 """ 

36 client = _get_gracedb_client(service_url=service_url) 

37 response = client.files(gracedb_id, filename) 

38 return response.json() 

39 

40 

41def get_gracedb_data( 

42 gracedb_id, superevent=False, info=None, json=None, 

43 service_url=conf.gracedb_server 

44): 

45 """Grab data from GraceDB for a specific event. 

46 

47 Parameters 

48 ---------- 

49 gracedb_id: str 

50 the GraceDB id of the event you wish to retrieve the data for 

51 superevent: Bool, optional 

52 True if the gracedb_id you are providing is a superevent 

53 info: str/list, optional 

54 either a string or list of strings for information you wish to 

55 retrieve 

56 json: dict, optional 

57 data that you have already downloaded from gracedb 

58 service_url: str, optional 

59 service url you wish to use when accessing data from GraceDB 

60 """ 

61 client = _get_gracedb_client(service_url=service_url) 

62 if json is None and superevent: 

63 json = client.superevent(gracedb_id).json() 

64 elif json is None: 

65 try: 

66 json = client.superevent(gracedb_id).json() 

67 except HTTPError: 

68 json = client.event(gracedb_id).json() 

69 

70 if isinstance(info, str) and info in json.keys(): 

71 return str(json[info]) 

72 elif isinstance(info, str): 

73 raise AttributeError( 

74 "Could not find '{}' in the gracedb dictionary. Available entries " 

75 "are: {}".format(info, ", ".join(json.keys())) 

76 ) 

77 elif isinstance(info, list): 

78 data = {} 

79 for _info in info: 

80 if _info in json.keys(): 

81 data[_info] = json[_info] 

82 else: 

83 logger.warning( 

84 "Unable to find any information for '{}'".format(_info) 

85 ) 

86 return data 

87 elif info is None: 

88 return json 

89 else: 

90 raise ValueError( 

91 "info data not understood. Please provide either a list or string" 

92 )