Merge pull request #35 from lsst-sitcom/tickets/DM-43209

DM-43209: Add type annotations

.pre-commit-config.yaml

@@ -12,7 +12,7 @@ repos:
         additional_dependencies: [toml]
 
   - repo: https://github.com/psf/black
-    rev: 23.12.1
+    rev: 24.1.1
     hooks:
       - id: black
 

python/lsst/summit/extras/animation.py

@@ -26,11 +26,14 @@
 import shutil
 import subprocess
 import uuid
+from typing import Any
 
 import matplotlib.pyplot as plt
 
 import lsst.afw.display as afwDisplay
+import lsst.afw.image as afwImage
 import lsst.afw.math as afwMath
+import lsst.daf.butler as dafButler
 import lsst.meas.algorithms as measAlg
 from lsst.atmospec.utils import airMassFromRawMetadata, getTargetCentroidFromWcs
 from lsst.pipe.tasks.quickFrameMeasurement import QuickFrameMeasurementTask, QuickFrameMeasurementTaskConfig
@@ -55,19 +58,19 @@ class Animator:
 
     def __init__(
         self,
-        butler,
-        dataIdList,
-        outputPath,
-        outputFilename,
+        butler: dafButler.Butler,
+        dataIdList: list[dict],
+        outputPath: str,
+        outputFilename: str,
         *,
-        remakePngs=False,
-        clobberVideoAndGif=False,
-        keepIntermediateGif=False,
-        smoothImages=True,
-        plotObjectCentroids=True,
-        useQfmForCentroids=False,
-        dataProductToPlot="calexp",
-        debug=False,
+        remakePngs: bool = False,
+        clobberVideoAndGif: bool = False,
+        keepIntermediateGif: bool = False,
+        smoothImages: bool = True,
+        plotObjectCentroids: bool = True,
+        useQfmForCentroids: bool = False,
+        dataProductToPlot: str = "calexp",
+        debug: bool = False,
     ):
         self.butler = butler
         self.dataIdList = dataIdList
@@ -100,11 +103,12 @@ def __init__(
         self.disp.setImageColormap("gray")
         self.disp.scale("asinh", "zscale")
 
-        self.pngsToMakeDataIds = []
+        self.pngsToMakeDataIds: list[dict] = []
+
         self.preRun()  # sets the above list
 
     @staticmethod
-    def _strDataId(dataId):
+    def _strDataId(dataId: dict) -> str:
         """Make a dataId into a string suitable for use as a filename.
 
         Parameters
@@ -130,7 +134,7 @@ def _strDataId(dataId):
         dIdStr = dIdStr.replace(",", "-")
         return dIdStr
 
-    def dataIdToFilename(self, dataId, includeNumber=False, imNum=None):
+    def dataIdToFilename(self, dataId: dict, includeNumber: bool = False, imNum: int | None = None) -> str:
         """Convert dataId to filename.
 
         Returns a full path+filename by default. if includeNumber then
@@ -147,16 +151,16 @@ def dataIdToFilename(self, dataId, includeNumber=False, imNum=None):
             filename = self.basicTemplate % (dIdStr, self.dataProductToPlot)
             return os.path.join(self.pngPath, filename)
 
-    def exists(self, obj):
+    def exists(self, obj: Any) -> bool:
         if isinstance(obj, str):
             return os.path.exists(obj)
         raise RuntimeError("Other type checks not yet implemented")
 
-    def preRun(self):
+    def preRun(self) -> None:
         # check the paths work
         if not os.path.exists(self.pngPath):
             os.makedirs(self.pngPath)
-        assert os.path.exists(self.pngPath), f"Failed to create output dir: {self.pngsPath}"
+        assert os.path.exists(self.pngPath), f"Failed to create output dir: {self.pngPath}"
 
         if self.exists(self.outputFilename):
             if self.clobberVideoAndGif:
@@ -202,7 +206,7 @@ def preRun(self):
             msg += " because remakePngs=True"
         logger.info(msg)
 
-    def run(self):
+    def run(self) -> str | None:
         # make the missing pngs
         if self.pngsToMakeDataIds:
             logger.info("Creating necessary pngs...")
@@ -213,7 +217,7 @@ def run(self):
         # stage files in temp dir with numbers prepended to filenames
         if not self.dataIdList:
             logger.warning("No files to animate - nothing to do")
-            return
+            return None
 
         logger.info("Copying files to ordered temp dir...")
         pngFilesOriginal = [self.dataIdToFilename(d) for d in self.dataIdList]
@@ -248,7 +252,7 @@ def run(self):
         logger.info(f"Finished! Output at {self.outputFilename}")
         return self.outputFilename
 
-    def _titleFromExp(self, exp, dataId):
+    def _titleFromExp(self, exp: afwImage.Exposure, dataId: dict) -> str:
         expRecord = getExpRecordFromDataId(self.butler, dataId)
         obj = expRecord.target_name
         expTime = expRecord.exposure_time
@@ -266,7 +270,9 @@ def _titleFromExp(self, exp, dataId):
         title += f"Object: {obj} expTime: {expTime}s Filter: {filt} Grating: {grating} Airmass: {airmass:.3f}"
         return title
 
-    def getStarPixCoord(self, exp, doMotionCorrection=True, useQfm=False):
+    def getStarPixCoord(
+        self, exp: Any, doMotionCorrection: bool = True, useQfm: bool = False
+    ) -> tuple[float, float] | None:
         target = exp.visitInfo.object
 
         if self.useQfmForCentroids:
@@ -281,7 +287,7 @@ def getStarPixCoord(self, exp, doMotionCorrection=True, useQfm=False):
             pixCoord = getTargetCentroidFromWcs(exp, target, doMotionCorrection=doMotionCorrection)
         return pixCoord
 
-    def makePng(self, dataId, saveFilename):
+    def makePng(self, dataId: dict, saveFilename: str) -> None:
         if self.exists(saveFilename) and not self.remakePngs:  # should not be possible due to prerun
             assert False, f"Almost overwrote {saveFilename} - how is this possible?"
 
@@ -332,7 +338,7 @@ def makePng(self, dataId, saveFilename):
         del exp
         gc.collect()
 
-    def pngsToMp4(self, indir, outfile, framerate, verbose=False):
+    def pngsToMp4(self, indir: str, outfile: str, framerate: float, verbose: bool = False) -> None:
         """Create the movie with ffmpeg, from files."""
         # NOTE: the order of ffmpeg arguments *REALLY MATTERS*.
         # Reorder them at your own peril!
@@ -370,11 +376,11 @@ def pngsToMp4(self, indir, outfile, framerate, verbose=False):
 
         subprocess.check_call(r" ".join(cmd), shell=True)
 
-    def tidyUp(self, tempDir):
+    def tidyUp(self, tempDir: str) -> None:
         shutil.rmtree(tempDir)
         return
 
-    def _smoothExp(self, exp, smoothing, kernelSize=7):
+    def _smoothExp(self, exp: afwImage.Exposure, smoothing: float, kernelSize: int = 7) -> afwImage.Exposure:
         """Use for DISPLAY ONLY!
 
         Return a smoothed copy of the exposure
@@ -389,7 +395,9 @@ def _smoothExp(self, exp, smoothing, kernelSize=7):
         return newExp
 
 
-def animateDay(butler, dayObs, outputPath, dataProductToPlot="quickLookExp"):
+def animateDay(
+    butler: dafButler.Butler, dayObs: int, outputPath: str, dataProductToPlot: str = "quickLookExp"
+) -> str | None:
     outputFilename = f"{dayObs}.mp4"
 
     onSkyIds = getLatissOnSkyDataIds(butler, startDate=dayObs, endDate=dayObs)

python/lsst/summit/extras/annotations.py

@@ -22,7 +22,7 @@
 from lsst.summit.extras.imageSorter import TAGS, ImageSorter
 
 
-def _idTrans(dataIdDictOrTuple):
+def _idTrans(dataIdDictOrTuple: dict | tuple[int, int]) -> tuple[int, int]:
     """Take a dataId and turn it into the internal tuple format."""
     if isinstance(dataIdDictOrTuple, tuple):
         return dataIdDictOrTuple
@@ -37,49 +37,49 @@ class Annotations:
     imageSorter.
     """
 
-    def __init__(self, filename):
+    def __init__(self, filename: str):
         self.filename = filename
         self.tags, self.notes = self._load(filename)
 
-    def _load(self, filename):
+    def _load(self, filename: str) -> tuple[dict, dict]:
         """Load tags and notes from specified file."""
         tags, notes = ImageSorter.loadAnnotations(filename)
         return tags, notes
 
-    def getTags(self, dataId):
+    def getTags(self, dataId: dict | tuple[int, int]) -> str | None:
         """Get the tags for a specified dataId.
 
         Empty string means no tags, None means not examined"""
         return self.tags.get(_idTrans(dataId), None)
 
-    def getNotes(self, dataId):
+    def getNotes(self, dataId: dict | tuple[int, int]) -> str | None:
         """Get the notes for the specified dataId."""
         return self.notes.get(_idTrans(dataId), None)
 
-    def hasTags(self, dataId, flags):
+    def hasTags(self, dataId: dict | tuple[int, int], flags: str) -> bool | None:
         """Check if a dataId has all the specificed tags"""
         tag = self.getTags(dataId)
         if tag is None:  # not just 'if tag' becuase '' is not the same as None but both as False-y
             return None
         return all(i in tag for i in flags.upper())
 
-    def getListOfCheckedData(self):
+    def getListOfCheckedData(self) -> list:
         """Return a list of all dataIds which have been examined."""
         return sorted(list(self.tags.keys()))
 
-    def getListOfDataWithNotes(self):
+    def getListOfDataWithNotes(self) -> list:
         """Return a list of all dataIds which have notes associated."""
         return sorted(list(self.notes.keys()))
 
-    def isExamined(self, dataId):
+    def isExamined(self, dataId: dict) -> bool:
         """Check if the dataId has been examined or not."""
         return _idTrans(dataId) in self.tags
 
-    def printTags(self):
+    def printTags(self) -> None:
         """Display the list of tag definitions."""
         print(TAGS)
 
-    def getIdsWithGivenTags(self, tags, exactMatches=False):
+    def getIdsWithGivenTags(self, tags: str, exactMatches: bool = False) -> list:
         if exactMatches:
             return [
                 dId

python/lsst/summit/extras/assessQFM.py

@@ -26,6 +26,7 @@
 import numpy as np
 import pandas as pd
 
+import lsst.daf.butler as dafButler
 import lsst.summit.utils.butlerUtils as butlerUtils
 from lsst.pipe.tasks.quickFrameMeasurement import QuickFrameMeasurementTask, QuickFrameMeasurementTaskConfig
 
@@ -61,13 +62,13 @@ class AssessQFM:
 
     def __init__(
         self,
-        butler,
-        dataProduct="quickLookExp",
-        dataset="data/qfm_baseline_assessment.parq",
-        successCut=2,
-        nearSuccessCut=10,
-        donutCut=10,
-        logLevel=50,
+        butler: dafButler.Butler,
+        dataProduct: str = "quickLookExp",
+        dataset: str = "data/qfm_baseline_assessment.parq",
+        successCut: int = 2,
+        nearSuccessCut: int = 10,
+        donutCut: int = 10,
+        logLevel: int = 50,
     ):
         self.butler = butler
 
@@ -99,7 +100,7 @@ def __init__(
             "U": "Ambiguous",  # Centroid is on a star, but unclear whether it is the brightest
         }
 
-    def run(self, nSamples=None, nProcesses=1, outputFile=None):
+    def run(self, nSamples: int | None = None, nProcesses: int = 1, outputFile: str | None = None) -> None:
         """Run quickFrameMeasurement on a sample dataset, save the new results,
         and compare them with the baseline, vetted by-eye results.
 
@@ -136,7 +137,7 @@ def run(self, nSamples=None, nProcesses=1, outputFile=None):
 
         self.compareToBaseline(qfmResults)
 
-    def _runQFM(self, testset):
+    def _runQFM(self, testset: pd.DataFrame) -> pd.DataFrame:
         """Run quickFrameMeasurement on a subset of the dataset.
 
         Parameters
@@ -176,7 +177,7 @@ def _runQFM(self, testset):
                 qfmRes["finalTag"] = "F"
         return qfmResults
 
-    def compareToBaseline(self, comparisonData):
+    def compareToBaseline(self, comparisonData: pd.DataFrame) -> None:
         """Compare a table of quickFrameMeasurement results with the
         baseline vetted data, and print output of the comparison.