From 542b13426cb66ad97afdcd793d89c919e32740f2 Mon Sep 17 00:00:00 2001
From: Constantine Evans <constantine@costi.net>
Date: Tue, 17 Dec 2024 01:14:06 -0500
Subject: [PATCH] mypy type checking fixes

---
 src/qslib/experiment.py         |  4 ++--
 src/qslib/machine.py            | 30 +++++++++++++++++++++++++-----
 src/qslib/plate_setup.py        |  2 +-
 src/qslib/qsconnection_async.py | 10 ++++++----
 4 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/src/qslib/experiment.py b/src/qslib/experiment.py
index 66d8dde..7947722 100644
--- a/src/qslib/experiment.py
+++ b/src/qslib/experiment.py
@@ -831,7 +831,7 @@ def sync_from_machine(
                 )
                 log.debug(f"checking {f['path']} mtime {f['mtime']} to {sdspath}")
                 if os.path.exists(sdspath) and os.path.getmtime(sdspath) >= float(
-                    f["mtime"]
+                    f["mtime"].timestamp()
                 ):
                     log.debug(f"{sdspath} has {os.path.getmtime(sdspath)}")
                     continue
@@ -878,7 +878,7 @@ def sync_from_machine(
                                 b.write(machine.read_file(f["path"]))
                 else:
                     raise NotImplementedError
-                os.utime(sdspath, (f["atime"], f["mtime"]))
+                os.utime(sdspath, (f["atime"].timestamp(), f["mtime"].timestamp()))
 
             # The message log is tricky. Ideally we'd use rsync or wc+tail. TODO
             self._update_from_files()
diff --git a/src/qslib/machine.py b/src/qslib/machine.py
index c731f2c..f8adc92 100644
--- a/src/qslib/machine.py
+++ b/src/qslib/machine.py
@@ -338,9 +338,17 @@ def list_files(
         leaf: str = "FILE",
         verbose: Literal[False] = False,
         recursive: bool = False,
-    ) -> list[str] | list[FileListInfo]: ...
-
+    ) -> list[str]: ...
 
+    @overload
+    def list_files(
+        self,
+        path: str,
+        *,
+        leaf: str = "FILE",
+        verbose: bool = False,
+        recursive: bool = False,
+    ) -> list[str] | list[FileListInfo]: ...
 
     @_ensure_connection(AccessLevel.Observer)
     def list_files(
@@ -394,8 +402,19 @@ def write_file(self, path: str, data: str | bytes) -> None:
             + b"\n</quote.base64>"
         )
 
+    @overload
+    def list_runs_in_storage(self, glob: str = "*", *, verbose: Literal[True]) -> list[FileListInfo]: ...
+
+    @overload
+    def list_runs_in_storage(self, glob: str = "*", *, verbose: Literal[False] = False) -> list[str]: ...
+
+    @overload
+    def list_runs_in_storage(self, glob: str = "*", *, verbose: bool = False) -> list[str] | list[FileListInfo]: ...
+
     @_ensure_connection(AccessLevel.Observer)
-    def list_runs_in_storage(self, glob: str = "*", verbose: bool = False) -> list[str]:
+    def list_runs_in_storage(
+        self, glob: str = "*", *, verbose: bool = False
+    ) -> list[str] | list[FileListInfo]:
         """List runs in machine storage.
 
         Returns
@@ -407,12 +426,13 @@ def list_runs_in_storage(self, glob: str = "*", verbose: bool = False) -> list[s
         """
         if not glob.endswith("eds"):
             glob = f"{glob}eds"
-        a = self.list_files(f"public_run_complete:{glob}", verbose=verbose)
         if not verbose:
             return [
-                re.sub("^public_run_complete:", "", s)[:-4] for s in a
+                re.sub("^public_run_complete:", "", s)[:-4]
+                for s in self.list_files(f"public_run_complete:{glob}", verbose=False)
             ]
         else:
+            a = self.list_files(f"public_run_complete:{glob}", verbose=True)
             for e in a:
                 e["path"] = re.sub("^public_run_complete:", "", e["path"])[:-4]
             return a
diff --git a/src/qslib/plate_setup.py b/src/qslib/plate_setup.py
index 14f0531..667f8d9 100644
--- a/src/qslib/plate_setup.py
+++ b/src/qslib/plate_setup.py
@@ -32,7 +32,7 @@
 
 if TYPE_CHECKING:
     from kithairon import PickList, Labware
-    from typing import Self
+    from typing_extensions import Self
 
 from .qsconnection_async import QSConnectionAsync
 
diff --git a/src/qslib/qsconnection_async.py b/src/qslib/qsconnection_async.py
index b3c5288..cea4119 100644
--- a/src/qslib/qsconnection_async.py
+++ b/src/qslib/qsconnection_async.py
@@ -24,7 +24,7 @@
 from .qs_is_protocol import CommandError, Error, NoMatch, QS_IS_Protocol
 from .scpi_commands import AccessLevel, ArgList, SCPICommand
 
-class FileListInfo(TypedDict):
+class FileListInfo(TypedDict, total=False):
     """Information about a file when verbose=True"""
     path: str
     type: str
@@ -32,6 +32,8 @@ class FileListInfo(TypedDict):
     mtime: datetime
     atime: datetime
     ctime: datetime
+    state: str
+    collected: bool
 
 log = logging.getLogger(__name__)
 
@@ -180,8 +182,8 @@ async def list_files(
                 )
                 if rm is None:
                     ag = ArgList.from_string(x)
-                    d: dict[str, str | float | int] = {}
-                    d["path"] = ag.args[0]
+                    d: dict[str, Any] = {}
+                    d["path"] = cast(str, ag.args[0])
                     d |= ag.opts
                 else:
                     d = {}
@@ -196,7 +198,7 @@ async def list_files(
                         cast(str, d["path"]), leaf=leaf, verbose=True, recursive=True
                     )
                 else:
-                    ret.append(d)
+                    ret.append(cast(FileListInfo, d))
             return ret
 
     async def compile_eds(self, run_name: str) -> None: