From 452f58c5d759db0fc94e057c3f2579c00eb77dd9 Mon Sep 17 00:00:00 2001 From: Tyler Pritchard Date: Fri, 2 Feb 2024 12:38:25 -0500 Subject: [PATCH] updated tests & bugfixes assosciated with them --- src/newlk_search/search.py | 16 ++++++++-------- tests/test_search.py | 14 +++++++------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/newlk_search/search.py b/src/newlk_search/search.py index a908bc1..d7d4908 100644 --- a/src/newlk_search/search.py +++ b/src/newlk_search/search.py @@ -661,7 +661,7 @@ def search_cubedata( exptime: Union[str, int, tuple] = None, cadence: Union[str, int, tuple] = None, mission: Union[str, tuple] = ("Kepler", "K2", "TESS"), - product: Union[str, tuple[str]] = ("Target Pixel","ffi"), + filetype: Union[str, list[str]] = ("Target Pixel","ffi"), author: Union[str, tuple] = None, quarter: Union[int, list[int]] = None, month: Union[int, list[int]] = None, @@ -674,7 +674,7 @@ def search_cubedata( return _search_products( target, radius=radius, - filetype=product, + filetype=filetype, exptime=exptime or cadence, mission=mission, provenance_name=author, @@ -929,7 +929,7 @@ def _search_products( # We will handle them separately using something else (tesswcs? tesspoint?) # Also the above is written like FFI kepler products exist, but I don't think they do? - extra_query_criteria["dataproduct_type"] = ["cube", "timeseries"] + #extra_query_criteria["dataproduct_type"] = ["cube", "timeseries"] # Query Mast to get a list of observations @@ -1005,7 +1005,7 @@ def _search_products( if f"c{row['sequence_number']}{half}" in row["productFilename"]: seq_num[index] = f"{int(row['sequence_number']):02d}{letter}" - joint_table["mission"] = [f" {proj} {obs_prefix.get(pref, '')} {seq}" + joint_table["mission"] = [f"{proj} {obs_prefix.get(pref, '')} {seq}" for proj, pref, seq in zip( joint_table['project'], joint_table['project'].values.astype(str), @@ -1023,7 +1023,6 @@ def _search_products( month=month, limit=limit, ) - print(masked_result) log.debug(f"MAST found {len(masked_result)} matching data products.") #masked_result["distance"].info.format = ".1f" # display <0.1 arcsec @@ -1273,18 +1272,19 @@ def _filter_products( mask |= _mask_kepler_products(products, quarter=quarter, month=month) # HLSP products need to be filtered by extension - if filetype.lower() == "lightcurve": + if "lightcurve" in [x.lower() for x in np.atleast_1d(filetype)]: mask &= np.array( [uri.lower().endswith("lc.fits") for uri in products["productFilename"]] ) - elif filetype.lower() == "target pixel": + # TODO:The elifs only allow for 1 type (target pixel or ffi), is that the behavior we want? + elif "target pixel" in [x.lower() for x in np.atleast_1d(filetype)]: mask &= np.array( [ uri.lower().endswith(("tp.fits", "targ.fits.gz")) for uri in products["productFilename"] ] ) - elif filetype.lower() == "ffi": + elif "ffi" in [x.lower() for x in np.atleast_1d(filetype)]:: mask &= np.array(["TESScut" in desc for desc in products["description"]]) # Allow only fits files diff --git a/tests/test_search.py b/tests/test_search.py index 41e2689..631156a 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -25,9 +25,9 @@ from lightkurve.utils import LightkurveWarning, LightkurveError from src.newlk_search.search import ( - search_lightcurve, - search_targetpixelfile, - search_tesscut, + search_timeseries, + search_cubedata, + #search_tesscut, SearchError, SearchResult, log, @@ -58,7 +58,7 @@ def remove_custom_config(): #@pytest.mark.remote_data def test_search_targetpixelfile(): # EPIC 210634047 was observed twice in long cadence - assert len(search_targetpixelfile("EPIC 210634047", mission="K2").table) == 2 + assert len(search_cubedata("EPIC 210634047", mission="K2").table) == 2 # ...including Campaign 4 assert ( len(search_targetpixelfile("EPIC 210634047", mission="K2", campaign=4).table) @@ -551,7 +551,7 @@ def test_qlp_ffi_lightcurve(): search = search_lightcurve("TrES-2b", sector=26, author="qlp") assert len(search) == 1 assert search.author[0] == "QLP" - assert search.exptime[0] == 30 * u.minute # Sector 26 had 30-minute FFIs + assert search.exptime[0] == 1800 * u.second # Sector 26 had 30-minute FFIs lc = search.download() all(lc.flux == lc.kspsap_flux) @@ -562,7 +562,7 @@ def test_spoc_ffi_lightcurve(): search = search_lightcurve("TrES-2b", sector=26, author="tess-spoc") assert len(search) == 1 assert search.author[0] == "TESS-SPOC" - assert search.exptime[0] == 30 * u.minute # Sector 26 had 30-minute FFIs + assert search.exptime[0] == 1800 * u.second # Sector 26 had 30-minute FFIs lc = search.download() all(lc.flux == lc.pdcsap_flux) @@ -630,7 +630,7 @@ def test_customize_search_result_display_case_nonexistent_column(): # One typical case is that some columns are in the result of # search_lightcurve() / search_targetpixelfile(), but not in those of search_tesscut() - search = search_lightcurve("TIC390021728") + search = search_timeseries("TIC390021728") search.display_extra_columns = ['foo_col'] assert 'foo_col' not in search.__repr__()