move gwosc segment query logic from gstlal_query_gwosc_segments to segments.query_gwosc_segments

gstlal-ugly/bin/gstlal_query_gwosc_segments

@@ -19,16 +19,14 @@
 
 import argparse
 import sys
-import ssl
-import urllib.request
-import warnings
 
+from lal import LIGOTimeGPS
 from ligo import segments
 from ligo.lw import ligolw
-from ligo.lw import utils as ligolw_utils
 from ligo.lw.utils import process as ligolw_process
 from ligo.lw.utils import segments as ligolw_segments
-from lal import LIGOTimeGPS
+
+from gstlal.segments import query_gwosc_segments
 
 #
 # Define/parse command line options
@@ -49,44 +47,22 @@ for x in args.ifo:
 	instruments |= set([x[i*2:i*2+2] for i in range(int(len(x) / 2))])
 
 #
-# determine GWOSC URL to query from
-#
-
-urlbase = "https://gw-openscience.org/timeline/segments"
-if span[0] in segments.segment(1126051217, 1137254417):
-	query_url = f"{urlbase}/O1"
-elif span[0] in segments.segment(1164556817, 1187733618):
-	query_url = f"{urlbase}/O2_16KHZ_R1"
-elif span[0] in segments.segment(1238166018, 1253977218):
-	query_url = f"{urlbase}/O3a_16KHZ_R1"
-else:
-	raise ValueError("GPS times requested not in GWOSC")
-
-#
-# Set up SSL context
+# Query segments
 #
 
-context = ssl.create_default_context()
-if args.no_cert:
-	context.check_hostname = False
-	context.verify_mode = ssl.CERT_NONE
+print("Downloading segment lists ...", file=sys.stderr)
+segments = query_gwosc_segments(instruments, args.start, args.end, verify_certs=(not args.no_cert))
 
 #
-# Retrieve segments
+# Write segments to disk
 #
 
 xmldoc = ligolw.Document()
 xmldoc.appendChild(ligolw.LIGO_LW())
-process = ligolw_process.register_to_xmldoc(xmldoc, "gstlal_query_gwosc_segments", {"": args.start, "": args.end})
-with ligolw_segments.LigolwSegments(xmldoc, process) as segdoc:
-	print("Downloading segment lists ...", file=sys.stderr)
-	for instrument in instruments:
-		url = f"{query_url}/{instrument}_DATA/{span[0]}/{abs(span)}"
-		urldata = urllib.request.urlopen(url, context=context).read().decode('utf-8')
-		with warnings.catch_warnings():
-			warnings.filterwarnings("ignore", category=FutureWarning)
-			segdoc.insert_from_segwizard(urldata.splitlines(), set((instrument,)), "datasegments")
-	segdoc.coalesce()
+process_params = {"start": args.start, "end": args.end}
+process = ligolw_process.register_to_xmldoc(xmldoc, sys.argv[0], process_params)
+with ligolw_segments.LigolwSegments(xmldoc, process) as lwseglists:
+	lwseglists.insert_from_segmentlistdict(segments, "datasegments")
 process.set_end_time_now()
+
 ligolw_utils.write_filename(xmldoc, args.output, verbose=True)
-xmldoc.unlink()

gstlal/python/segments.py

@@ -17,9 +17,75 @@
 
 import itertools
 import math
-from typing import Iterable
+import ssl
+import urllib.request
+from typing import Iterable, Union
+import warnings
 
 from ligo import segments
+from ligo.lw import ligolw
+from ligo.lw.utils import process as ligolw_process
+from ligo.lw.utils import segments as ligolw_segments
+from lal import LIGOTimeGPS
+
+
+def query_gwosc_segments(
+	instruments: Union[str, Iterable],
+	start: Union[int, LIGOTimeGPS],
+	end: Union[int, LIGOTimeGPS],
+	verify_certs: bool = True,
+) -> segments.segmentlistdict:
+	"""Query GWOSC for science segments.
+
+	Args:
+		instruments:
+			Union[str, Iterable], the instruments to query segments for
+		start:
+			Union[int, LIGOTimeGPS], the GPS start time
+		end:
+			Union[int, LIGOTimeGPS], the GPS end time
+		verify_certs:
+			bool, default True, whether to verify SSL certificates when querying GWOSC.
+
+	Returns:
+		segmentlistdict, the queried segments
+
+	"""
+	span = segments.segment(LIGOTimeGPS(start), LIGOTimeGPS(end))
+	if isinstance(instruments, str):
+		instruments = [instruments]
+
+	# determine GWOSC URL to query from
+	urlbase = "https://gw-openscience.org/timeline/segments"
+	if span[0] in segments.segment(1126051217, 1137254417):
+		query_url = f"{urlbase}/O1"
+	elif span[0] in segments.segment(1164556817, 1187733618):
+		query_url = f"{urlbase}/O2_16KHZ_R1"
+	elif span[0] in segments.segment(1238166018, 1253977218):
+		query_url = f"{urlbase}/O3a_16KHZ_R1"
+	else:
+		raise ValueError("GPS times requested not in GWOSC")
+
+	# Set up SSL context
+	context = ssl.create_default_context()
+	if not verify_certs:
+		context.check_hostname = False
+		context.verify_mode = ssl.CERT_NONE
+
+	# Retrieve segments
+	xmldoc = ligolw.Document()
+	xmldoc.appendChild(ligolw.LIGO_LW())
+	process = ligolw_process.register_to_xmldoc(xmldoc, "gstlal.segments", {})
+	with ligolw_segments.LigolwSegments(xmldoc, process) as segdoc:
+		for instrument in instruments:
+			url = f"{query_url}/{instrument}_DATA/{span[0]}/{abs(span)}"
+			urldata = urllib.request.urlopen(url, context=context).read().decode('utf-8')
+			with warnings.catch_warnings():
+				warnings.filterwarnings("ignore", category=FutureWarning)
+				segdoc.insert_from_segwizard(urldata.splitlines(), set((instrument,)), "datasegments")
+		segdoc.coalesce()
+
+	return ligolw_segments.segmenttable_get_by_name(xmldoc, "datasegments")
 
 
 def analysis_segments(