From 077e5a3f77896076495405efc136a5df1b649785 Mon Sep 17 00:00:00 2001
From: "Brett M. Morris" <bmmorris@stsci.edu>
Date: Wed, 20 Dec 2023 09:59:45 -0500
Subject: [PATCH] adding kepler and tess specific TPF translators

---
 lcviz/parsers.py | 14 ++++++++------
 lcviz/utils.py   | 22 +++++++++++++++++-----
 2 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/lcviz/parsers.py b/lcviz/parsers.py
index 62d35586..047af41f 100644
--- a/lcviz/parsers.py
+++ b/lcviz/parsers.py
@@ -10,6 +10,13 @@
 def light_curve_parser(app, file_obj, data_label=None, show_in_viewer=True, **kwargs):
     time_viewer_reference_name = app._jdaviz_helper._default_time_viewer_reference_name
 
+    # load a LightCurve or TargetPixelFile object:
+    cls_with_translator = (
+        lightkurve.LightCurve,
+        lightkurve.targetpixelfile.KeplerTargetPixelFile,
+        lightkurve.targetpixelfile.TessTargetPixelFile
+    )
+
     # load local FITS file from disk by its path:
     if isinstance(file_obj, str) and os.path.exists(file_obj):
         if data_label is None:
@@ -18,12 +25,7 @@ def light_curve_parser(app, file_obj, data_label=None, show_in_viewer=True, **kw
         # read the light curve:
         light_curve = lightkurve.read(file_obj)
 
-    # load a LightCurve object:
-    elif isinstance(file_obj, lightkurve.LightCurve):
-        light_curve = file_obj
-
-    # load a LightCurve object:
-    elif isinstance(file_obj, lightkurve.targetpixelfile.KeplerTargetPixelFile):
+    elif isinstance(file_obj, cls_with_translator):
         light_curve = file_obj
 
     # make a data label:
diff --git a/lcviz/utils.py b/lcviz/utils.py
index 7ede32cf..c6344a60 100644
--- a/lcviz/utils.py
+++ b/lcviz/utils.py
@@ -11,8 +11,10 @@
 from lightkurve import (
     LightCurve, KeplerLightCurve, TessLightCurve, FoldedLightCurve
 )
-from lightkurve.targetpixelfile import KeplerTargetPixelFile, TargetPixelFileFactory
-from lightkurve.utils import KeplerQualityFlags
+from lightkurve.targetpixelfile import (
+    KeplerTargetPixelFile, TessTargetPixelFile, TargetPixelFileFactory
+)
+from lightkurve.utils import KeplerQualityFlags, TessQualityFlags
 
 from astropy import units as u
 from astropy.table import QTable
@@ -287,8 +289,8 @@ def to_object(self, data_or_subset):
         return LightCurve(table, **kwargs)
 
 
-@data_translator(KeplerTargetPixelFile)
-class KeplerTPFHandler:
+class TPFHandler:
+    quality_flag_cls = None
     tpf_attrs = ['flux', 'flux_bkg', 'flux_bkg_err', 'flux_err']
     meta_attrs = [
         'cadenceno',
@@ -443,7 +445,7 @@ def to_object(self, data_or_subset):
             if hasattr(tpf, attr) and getattr(getattr(tpf, attr), 'fset', None) is not None:
                 setattr(tpf, attr, data.meta[attr])
 
-        tpf.quality_mask = KeplerQualityFlags.create_quality_mask(
+        tpf.quality_mask = self.quality_flag_cls.create_quality_mask(
             quality_array=tpf.hdu[1].data["QUALITY"], bitmask=tpf.quality_bitmask
         )
 
@@ -482,6 +484,16 @@ class TessLightCurveHandler(LightCurveHandler):
     pass
 
 
+@data_translator(KeplerTargetPixelFile)
+class KeplerTPFHandler(TPFHandler):
+    quality_flag_cls = KeplerQualityFlags
+
+
+@data_translator(TessTargetPixelFile)
+class TessTPFHandler(TPFHandler):
+    quality_flag_cls = TessQualityFlags
+
+
 # plugin component filters
 def data_not_folded(data):
     return data.meta.get('_LCVIZ_EPHEMERIS', None) is None