diff --git a/astrodb_utils/photometry.py b/astrodb_utils/photometry.py index eaad73f..7d8f8e2 100644 --- a/astrodb_utils/photometry.py +++ b/astrodb_utils/photometry.py @@ -109,7 +109,11 @@ def ingest_photometry( return flags # TODO: Make sure band exists in the PhotometryFilters table - band_match = db.query(db.PhotometryFilters).filter(db.PhotometryFilters.c.band == band).table() + band_match = ( + db.query(db.PhotometryFilters) + .filter(db.PhotometryFilters.c.band == band) + .table() + ) if len(band_match) == 0: msg = f"Band {band} not found in PhotometryFilters table." if raise_error: @@ -121,7 +125,11 @@ def ingest_photometry( # If telescope is provided, make sure it exists in the Telescopes table if telescope is not None: - telescope_match = db.query(db.Telescopes).filter(db.Telescopes.c.telescope == telescope).table() + telescope_match = ( + db.query(db.Telescopes) + .filter(db.Telescopes.c.telescope == telescope) + .table() + ) if len(telescope_match) == 0: msg = f"Telescope {telescope} not found in Telescopes table." if raise_error: @@ -146,7 +154,9 @@ def ingest_photometry( "source": db_name, "band": band, "regime": regime, - "magnitude": str(magnitude), # Convert to string to maintain significant digits + "magnitude": str( + magnitude + ), # Convert to string to maintain significant digits "magnitude_error": mag_error, "telescope": telescope, "epoch": epoch, @@ -186,12 +196,23 @@ def ingest_photometry( return flags -def ingest_photometry_filter(db, *, telescope=None, instrument=None, filter_name=None, ucd=None): +def ingest_photometry_filter( + db, + *, + telescope=None, + instrument=None, + filter_name=None, + ucd=None, + wavelength_col_name: str = "effective_wavelength_angstroms", + width_col_name: str = "width_angstroms", +): """ Add a new photometry filter to the database """ # Fetch existing telescopes, add if missing - existing = db.query(db.Telescopes).filter(db.Telescopes.c.telescope == telescope).table() + existing = ( + db.query(db.Telescopes).filter(db.Telescopes.c.telescope == telescope).table() + ) if len(existing) == 0: with db.engine.connect() as conn: conn.execute(db.Telescopes.insert().values({"telescope": telescope})) @@ -201,16 +222,24 @@ def ingest_photometry_filter(db, *, telescope=None, instrument=None, filter_name logger.info(f"Telescope {telescope} already exists.") # Fetch existing instruments, add if missing - existing = db.query(db.Instruments).filter(db.Instruments.c.instrument == instrument).table() + existing = ( + db.query(db.Instruments) + .filter(db.Instruments.c.instrument == instrument) + .table() + ) if len(existing) == 0: - ingest_instrument(db, telescope=telescope, instrument=instrument, mode="Imaging") + ingest_instrument( + db, telescope=telescope, instrument=instrument, mode="Imaging" + ) logger.info(f"Added instrument {instrument}.") else: logger.info(f"Instrument {instrument} already exists.") # Get data from SVO try: - filter_id, wave_eff, fwhm, width_effective = fetch_svo(telescope, instrument, filter_name) + filter_id, wave_eff, fwhm, width_effective = fetch_svo( + telescope, instrument, filter_name + ) logger.info( f"From SVO: Filter {filter_id} has effective wavelength {wave_eff} " f"and FWHM {fwhm} and width_effective {width_effective}." @@ -220,7 +249,6 @@ def ingest_photometry_filter(db, *, telescope=None, instrument=None, filter_name logger.error(msg) raise AstroDBError(msg) - if ucd is None: ucd = assign_ucd(wave_eff) logger.info(f"UCD for filter {filter_id} is {ucd}") @@ -233,8 +261,8 @@ def ingest_photometry_filter(db, *, telescope=None, instrument=None, filter_name { "band": filter_id, "ucd": ucd, - "effective_wavelength_angstroms": wave_eff.to(u.Angstrom).value, - "width_angstroms": width_effective.to(u.Angstrom).value, + wavelength_col_name: wave_eff.to(u.Angstrom).value, + width_col_name: width_effective.to(u.Angstrom).value, } ) ) diff --git a/tests/test_utils.py b/tests/test_utils.py index 0c80c51..2c6b347 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -21,8 +21,8 @@ def test_get_db_regime(db, input, db_regime): def test_get_db_regime_errors(db): # TESTS WHICH SHOULD FAIL with pytest.raises(AstroDBError) as error_message: - get_db_regime(db, "nir") - assert "Regime nir not found in database" in str(error_message.value) + get_db_regime(db, "not_a_regime") + assert "Regime not_a_regime not found in database" in str(error_message.value) with pytest.raises(AstroDBError) as error_message: get_db_regime(db, "xray")