Merge pull request #153 from jenshnielsen/mypy_fixes_12

Some typing improvements to qcodes dataset

plottr/apps/inspectr.py

@@ -17,7 +17,7 @@
 import sys
 import argparse
 import logging
-from typing import Optional, Sequence, List, Dict, Iterable, Union
+from typing import Optional, Sequence, List, Dict, Iterable, Union, cast
 from typing_extensions import TypedDict
 
 import pandas
@@ -146,8 +146,8 @@ def addRun(self, runId: int, **vals: str) -> None:
         lst.append(vals.get('experiment', ''))
         lst.append(vals.get('sample', ''))
         lst.append(vals.get('name', ''))
-        lst.append(vals.get('started date', '') + ' ' + vals.get('started time', ''))
-        lst.append(vals.get('completed date', '') + ' ' + vals.get('completed time', ''))
+        lst.append(vals.get('started_date', '') + ' ' + vals.get('started_time', ''))
+        lst.append(vals.get('completed_date', '') + ' ' + vals.get('completed_time', ''))
         lst.append(str(vals.get('records', '')))
         lst.append(vals.get('guid', ''))
 
@@ -424,7 +424,7 @@ def DBLoaded(self, dbdf: pandas.DataFrame) -> None:
     def updateDates(self) -> None:
         assert self.dbdf is not None
         if self.dbdf.size > 0:
-            dates = list(self.dbdf.groupby('started date').indices.keys())
+            dates = list(self.dbdf.groupby('started_date').indices.keys())
             self.dateList.updateDates(dates)
 
     ### reloading the db
@@ -456,7 +456,7 @@ def monitorTriggered(self) -> None:
     def setDateSelection(self, dates: Sequence[str]) -> None:
         if len(dates) > 0:
             assert self.dbdf is not None
-            selection = self.dbdf.loc[self.dbdf['started date'].isin(dates)].sort_index(ascending=False)
+            selection = self.dbdf.loc[self.dbdf['started_date'].isin(dates)].sort_index(ascending=False)
             self.runList.setRuns(selection.to_dict(orient='index'))
         else:
             self.runList.clear()
@@ -469,7 +469,8 @@ def setRunSelection(self, runId: int) -> None:
         if hasattr(ds, 'snapshot'):
             snap = ds.snapshot
 
-        structure = get_ds_structure(ds)
+        structure = cast(Dict[str, dict], get_ds_structure(ds))
+        # cast away typed dict so we can pop a key
         for k, v in structure.items():
             v.pop('values')
         contentInfo = {'Data structure': structure,

plottr/data/qcodes_dataset.py

@@ -8,6 +8,8 @@
 from operator import attrgetter
 from typing import Dict, List, Set, Union, TYPE_CHECKING, Any, Tuple, Optional, cast
 
+from typing_extensions import TypedDict
+
 import pandas as pd
 
 from qcodes.dataset.data_set import load_by_id
@@ -35,9 +37,38 @@ def _get_names_of_standalone_parameters(paramspecs: List['ParamSpec']
     return standalones
 
 
+class IndependentParameterDict(TypedDict):
+    unit: str
+    label: str
+    values: List[Any]
+
+
+class DependentParameterDict(IndependentParameterDict):
+    axes: List[str]
+
+
+DataSetStructureDict = Dict[str, Union[IndependentParameterDict, DependentParameterDict]]
+
+
+class DataSetInfoDict(TypedDict):
+    experiment: str
+    sample: str
+    name: str
+    completed_date: str
+    completed_time: str
+    started_date: str
+    started_time: str
+    structure: Optional[DataSetStructureDict]
+    records: int
+    guid: str
+
+
 # Tools for extracting information on runs in a database
 
-def get_ds_structure(ds: 'DataSet') -> Dict[str, Any]:
+
+def get_ds_structure(
+        ds: 'DataSet'
+) -> DataSetStructureDict:
     """
     Return the structure of the dataset, i.e., a dictionary in the form
         {
@@ -60,56 +91,68 @@ def get_ds_structure(ds: 'DataSet') -> Dict[str, Any]:
     in the returned structure.
     """
 
-    structure = {}
+    structure: DataSetStructureDict = {}
 
     paramspecs = ds.get_parameters()
 
     standalones = _get_names_of_standalone_parameters(paramspecs)
 
     for spec in paramspecs:
         if spec.name not in standalones:
-            structure[spec.name] = {'unit': spec.unit, 'label': spec.label, 'values': []}
             if len(spec.depends_on_) > 0:
-                structure[spec.name]['axes'] = list(spec.depends_on_)
-
+                structure[spec.name] = DependentParameterDict(unit=spec.unit,
+                                                              label=spec.label,
+                                                              values=[],
+                                                              axes=list(spec.depends_on_))
+            else:
+                structure[spec.name] = IndependentParameterDict(unit=spec.unit,
+                                                                label=spec.label,
+                                                                values=[])
     return structure
 
 
-def get_ds_info(ds: 'DataSet', get_structure: bool = True) -> Dict[str, Union[str,int, Dict[str, Any]]]:
+def get_ds_info(ds: 'DataSet', get_structure: bool = True) -> DataSetInfoDict:
     """
     Get some info on a DataSet in dict.
 
     if get_structure is True: return the datastructure in that dataset
     as well (key is `structure' then).
     """
-    ret: Dict[str, Union[str, int, Dict[str, Any]]] = {}
-    ret['experiment'] = ds.exp_name
-    ret['sample'] = ds.sample_name
-    ret['name'] = ds.name
-
     _complete_ts = ds.completed_timestamp()
     if _complete_ts is not None:
-        ret['completed date'] = _complete_ts[:10]
-        ret['completed time'] = _complete_ts[11:]
+        completed_date = _complete_ts[:10]
+        completed_time = _complete_ts[11:]
     else:
-        ret['completed date'] = ''
-        ret['completed time'] = ''
+        completed_date = ''
+        completed_time = ''
 
     _start_ts = ds.run_timestamp()
     if _start_ts is not None:
-        ret['started date'] = _start_ts[:10]
-        ret['started time'] = _start_ts[11:]
+        started_date = _start_ts[:10]
+        started_time = _start_ts[11:]
     else:
-        ret['started date'] = ''
-        ret['started time'] = ''
+        started_date = ''
+        started_time = ''
 
     if get_structure:
-        ret['structure'] = get_ds_structure(ds)
-
-    ret['records'] = ds.number_of_results
-    ret['guid'] = ds.guid
+        structure: Optional[DataSetStructureDict] = get_ds_structure(ds)
+    else:
+        structure = None
+
+    data = DataSetInfoDict(
+        experiment=ds.exp_name,
+        sample=ds.sample_name,
+        name=ds.name,
+        completed_date=completed_date,
+        completed_time=completed_time,
+        started_date=started_date,
+        started_time=started_time,
+        structure=structure,
+        records=ds.number_of_results,
+        guid=ds.guid
+    )
 
-    return ret
+    return data
 
 
 def load_dataset_from(path: str, run_id: int) -> 'DataSet':
@@ -126,7 +169,7 @@ def load_dataset_from(path: str, run_id: int) -> 'DataSet':
 
 def get_runs_from_db(path: str, start: int = 0,
                      stop: Union[None, int] = None,
-                     get_structure: bool = False) -> Dict[int, Dict[str, Any]]:
+                     get_structure: bool = False) -> Dict[int, DataSetInfoDict]:
     """
     Get a db ``overview`` dictionary from the db located in ``path``. The
     ``overview`` dictionary maps ``DataSet.run_id``s to dataset information as

plottr/utils/misc.py

@@ -3,10 +3,10 @@
 Various utility functions.
 """
 
-from typing import List, Tuple, TypeVar, Optional
+from typing import List, Tuple, TypeVar, Optional, Sequence
 
 
-def reorder_indices(lst: List, target: List) -> Tuple[int, ...]:
+def reorder_indices(lst: Sequence[str], target: Sequence[str]) -> Tuple[int, ...]:
     """
     Determine how to bring a list with unique entries to a different order.
 

test/pytest/test_qcodes_data.py

@@ -196,8 +196,8 @@ def test_get_ds_info(experiment):
 
         ds_info_with_empty_timestamps = get_ds_info(dataset,
                                                     get_structure=False)
-        assert ds_info_with_empty_timestamps['completed date'] == ''
-        assert ds_info_with_empty_timestamps['completed time'] == ''
+        assert ds_info_with_empty_timestamps['completed_date'] == ''
+        assert ds_info_with_empty_timestamps['completed_time'] == ''
 
     # timestamps are difficult to test for, so we will cheat here and
     # instead of hard-coding timestamps we will just get them from the dataset
@@ -208,11 +208,12 @@ def test_get_ds_info(experiment):
     expected_ds_info = {
         'experiment': '2d_softsweep',
         'sample': 'no sample',
-        'completed date': completed_ts[:10],
-        'completed time': completed_ts[11:],
-        'started date': started_ts[:10],
-        'started time': started_ts[11:],
+        'completed_date': completed_ts[:10],
+        'completed_time': completed_ts[11:],
+        'started_date': started_ts[:10],
+        'started_time': started_ts[11:],
         'name': 'results',
+        'structure': None,
         'records': 0,
         'guid': dataset.guid
     }