Address reviewer comments

detective/core.py

@@ -83,9 +83,8 @@ def fetch_entities(self) -> None:
         response = self.perform_query(query)
 
         # Parse the domains from the entities.
-        entities = [e[0] for e in response]
-        print(f"There are {len(entities)} entities with data")
-        self.entities = entities
+        self.entities = [e[0] for e in response]
+        print(f"There are {len(self.entities)} entities with data")
 
     def fetch_all_sensor_data(self, limit=50000) -> pd.DataFrame:
         """
@@ -98,8 +97,6 @@ def fetch_all_sensor_data(self, limit=50000) -> pd.DataFrame:
                 domain IN ('binary_sensor', 'sensor')
             AND
                 state NOT IN ('unknown', 'unavailable')
-            AND
-                last_changed = last_updated
             ORDER BY last_changed DESC
             LIMIT {limit}
             """

detective/functions.py

@@ -7,40 +7,23 @@
 UNKNOWN = "unknown"
 
 
-def format_binary_state(state: str):
-    """Return a binary for the state of binary sensors."""
-    if state == "on":
-        return 1
-    elif state == "off":
-        return 0
-    return state
-
-
 def get_device_class(attributes: dict):
     """Return the device class."""
-    device_class = attributes.get("device_class")
-    if device_class:
-        return device_class
-    return UNKNOWN
+    return attributes.get("device_class", UNKNOWN)
 
 
 def get_unit_of_measurement(attributes: dict):
     """Return the unit_of_measurement."""
-    unit_of_measurement = attributes.get("unit_of_measurement")
-    if unit_of_measurement:
-        return unit_of_measurement
-    return UNKNOWN
+    return attributes.get("unit_of_measurement", UNKNOWN)
 
 
 def get_friendly_name(attributes: dict):
     """Return the friendly_name."""
-    friendly_name = attributes.get("friendly_name")
-    if friendly_name:
-        return friendly_name
-    return UNKNOWN
+    return attributes.get("friendly_name", UNKNOWN)
 
 
 def generate_features(df: pd.DataFrame) -> pd.DataFrame:
+    """Generate features from the attributes."""
     df["attributes"] = df["attributes"].apply(json.loads)
     df["device_class"] = df["attributes"].apply(get_device_class)
     df["unit_of_measurement"] = df["attributes"].apply(get_unit_of_measurement)
@@ -49,10 +32,8 @@ def generate_features(df: pd.DataFrame) -> pd.DataFrame:
 
 
 def format_dataframe(df: pd.DataFrame) -> pd.DataFrame:
-    df["state"] = df["state"].apply(format_binary_state)
-    df["state"] = pd.to_numeric(
-        df["state"], errors="coerce"
-    )  # coerce will return NaN if unable to convert
+    """Convert states to numeric where possible and format the last_changed."""
+    df["state"] = pd.to_numeric(df["state"], errors="coerce")
     df["last_changed"] = pd.to_datetime(
         df["last_changed"].values, errors="ignore", utc=True
     ).tz_localize(None)