code tidy

Author: nickzoic

countess/gui/tabular.py

@@ -103,7 +103,7 @@ class TabularDataFrame(tk.Frame):
     columns: list[tk.Text] = []
     column_formats: list[str] = []
     scrollbar = None
-    sort_by_col = None
+    sort_by_col: Optional[int] = None
     sort_ascending = True
     callback: Optional[Callable[[int, int, bool], None]] = None
     click_callback: Optional[Callable[[int, int, int], None]] = None
@@ -198,6 +198,8 @@ def _column_click(self, col: int, ev):
 
     def refresh(self, new_offset: float = 0):
         # Refreshes the column widgets.
+        assert self.table
+        assert self.scrollbar
 
         self.offset = max(0, min(self.length - self.height, int(new_offset)))
 

countess/plugins/csv.py

@@ -3,7 +3,7 @@
 import gzip
 import logging
 from io import BufferedWriter, BytesIO
-from typing import Any, List, Optional, Sequence, Tuple, Union
+from typing import List, Optional, Sequence, Tuple, Union
 
 import duckdb
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
@@ -20,7 +20,7 @@
     StringParam,
 )
 from countess.core.plugins import DuckdbLoadFilePlugin, DuckdbSaveFilePlugin
-from countess.utils.duckdb import duckdb_escape_identifier, duckdb_escape_literal, duckdb_source_to_view
+from countess.utils.duckdb import duckdb_escape_identifier, duckdb_escape_literal
 from countess.utils.files import clean_filename
 
 CSV_FILE_TYPES: Sequence[Tuple[str, Union[str, List[str]]]] = [

countess/plugins/pivot.py

@@ -1,4 +1,3 @@
-import functools
 import logging
 from typing import Optional
 
@@ -7,12 +6,7 @@
 from countess import VERSION
 from countess.core.parameters import BooleanParam, ChoiceParam, PerColumnArrayParam
 from countess.core.plugins import DuckdbSimplePlugin
-from countess.utils.duckdb import (
-    duckdb_choose_special,
-    duckdb_escape_identifier,
-    duckdb_escape_literal,
-    duckdb_source_to_view,
-)
+from countess.utils.duckdb import duckdb_choose_special, duckdb_escape_identifier, duckdb_escape_literal
 
 logger = logging.getLogger(__name__)
 

countess/plugins/score.py

@@ -1,8 +1,7 @@
 from math import log
-from typing import Any, List, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
-import pandas as pd
 from scipy.optimize import curve_fit
 
 from countess import VERSION
@@ -64,7 +63,7 @@ def output_columns(self) -> dict[str, str]:
 
     def prepare(self, source):
         yaxis_prefix = self.columns.get_column_prefix()
-        suffix_set = set([k.removeprefix(yaxis_prefix) for k in source.columns if k.startswith(yaxis_prefix)])
+        suffix_set = {k.removeprefix(yaxis_prefix) for k in source.columns if k.startswith(yaxis_prefix)}
 
         if self.xaxis.is_not_none():
             xaxis_prefix = self.xaxis.get_column_prefix()
@@ -73,6 +72,8 @@ def prepare(self, source):
         self.suffixes = sorted(suffix_set)
 
     def transform(self, data: dict[str, Any]) -> Optional[dict[str, Any]]:
+        assert self.suffixes
+
         if self.xaxis.is_not_none():
             xaxis_prefix = self.xaxis.get_column_prefix()
             x_values = [data.get(xaxis_prefix + s) for s in self.suffixes]
@@ -85,17 +86,22 @@ def transform(self, data: dict[str, Any]) -> Optional[dict[str, Any]]:
             return None
 
         if self.log:
-            y_values = [log(y + 1) for y in y_values]
+            y_values = [log(y + 1) if y is not None else None for y in y_values]
         if self.normalize:
-            max_y = max(y_values)
-            y_values = [y / max_y for y in y_values]
+            max_y = max(y for y in y_values if y is not None)
+            y_values = [y / max_y if y is not None else None for y in y_values]
 
-        x_values, y_values = zip(*[(x, y) for x, y in zip(x_values, y_values) if x > 0 or y > 0])
+        x_values, y_values = zip(
+            *[(x, y) for x, y in zip(x_values, y_values) if x is not None and y is not None and (x > 0 or y > 0)]
+        )
         if len(x_values) < len(self.suffixes) / 2 + 1:
             return None
 
-        s, v = score(x_values, y_values)
-        if self.variance:
-            return {self.output.value: s, self.variance.value: v}
-        else:
-            return {self.output.value: s}
+        try:
+            s, v = score(x_values, y_values)
+            if self.variance:
+                return {self.output.value: s, self.variance.value: v}
+            else:
+                return {self.output.value: s}
+        except TypeError:
+            return None