rework dataset function and class structure

server/dataset/DatasetStats.py

@@ -1,50 +1,16 @@
 from collections import Counter, defaultdict
-import argparse, glob, json, re
+import argparse, json, re
 from datetime import datetime
 from pathlib import Path
 
+from server.dataset.DatasetIO import DatasetIO
+
 class DatasetStats:
   DAY_PATTERN = re.compile(r"(\d{4}-\d{2}-\d{2})")
 
   def __init__(self, input_files:list[str]):
     self.input_files = input_files
 
-  def _resolve_input_files(self):
-    resolved = []
-    seen = set()
-
-    for item in self.input_files:
-      path = Path(item)
-      if path.is_dir():
-        for file_path in sorted(path.rglob("*.jsonl")):
-          key = str(file_path.resolve())
-          if key in seen:
-            continue
-          seen.add(key)
-          resolved.append(file_path)
-        continue
-
-      if any(ch in item for ch in "*?[]"):
-        for match in sorted(glob.glob(item)):
-          file_path = Path(match)
-          if not file_path.is_file():
-            continue
-          key = str(file_path.resolve())
-          if key in seen:
-            continue
-          seen.add(key)
-          resolved.append(file_path)
-        continue
-
-      if path.is_file():
-        key = str(path.resolve())
-        if key in seen:
-          continue
-        seen.add(key)
-        resolved.append(path)
-
-    return resolved
-
   def _infer_day(self, file_path:Path):
     match = self.DAY_PATTERN.search(file_path.name)
     if match:
@@ -83,7 +49,7 @@ class DatasetStats:
     return None
 
   def analyze(self):
-    files = self._resolve_input_files()
+    files = DatasetIO.resolve_input_files(self.input_files)
 
     totals = {
       "rows": 0,
@@ -99,57 +65,52 @@ class DatasetStats:
 
     for file_path in files:
       day = self._infer_day(file_path)
-      with file_path.open("r", encoding="utf-8") as source:
-        for line in source:
-          if not line.strip():
-            continue
+      for row in DatasetIO.iter_jsonl_rows(file_path):
+        game_id = row.get("game_id")
+        if not game_id:
+          continue
 
-          row = json.loads(line)
-          game_id = row.get("game_id")
-          if not game_id:
-            continue
+        turn = int(row.get("turn", 0))
+        safe_options = self._extract_safe_options(row)
+        snake_type = row.get("snake_type", "unknown")
+        move = row.get("move", "unknown")
 
-          turn = int(row.get("turn", 0))
-          safe_options = self._extract_safe_options(row)
-          snake_type = row.get("snake_type", "unknown")
-          move = row.get("move", "unknown")
+        game_type = row.get("game_type", {})
+        if isinstance(game_type, dict):
+          game_type_name = game_type.get("name", "unknown")
+        else:
+          game_type_name = str(game_type)
 
-          game_type = row.get("game_type", {})
-          if isinstance(game_type, dict):
-            game_type_name = game_type.get("name", "unknown")
-          else:
-            game_type_name = str(game_type)
+        totals["rows"] += 1
+        totals["games"].add(game_id)
+        totals["snake_types"][snake_type] += 1
+        totals["game_types"][game_type_name] += 1
+        totals["moves"][move] += 1
+        totals["days"][day] += 1
 
-          totals["rows"] += 1
-          totals["games"].add(game_id)
-          totals["snake_types"][snake_type] += 1
-          totals["game_types"][game_type_name] += 1
-          totals["moves"][move] += 1
-          totals["days"][day] += 1
+        if game_id not in games:
+          games[game_id] = {
+            "game_id": game_id,
+            "day": day,
+            "snake_type": snake_type,
+            "game_type": game_type_name,
+            "rows": 0,
+            "max_turn": -1,
+            "safe_options_sum": 0,
+            "safe_options_count": 0,
+            "pressure_turns": 0,
+          }
 
-          if game_id not in games:
-            games[game_id] = {
-              "game_id": game_id,
-              "day": day,
-              "snake_type": snake_type,
-              "game_type": game_type_name,
-              "rows": 0,
-              "max_turn": -1,
-              "safe_options_sum": 0,
-              "safe_options_count": 0,
-              "pressure_turns": 0,
-            }
+        game = games[game_id]
+        game["rows"] += 1
+        game["max_turn"] = max(game["max_turn"], turn)
+        if isinstance(safe_options, int):
+          game["safe_options_sum"] += safe_options
+          game["safe_options_count"] += 1
+          if safe_options <= 2:
+            game["pressure_turns"] += 1
 
-          game = games[game_id]
-          game["rows"] += 1
-          game["max_turn"] = max(game["max_turn"], turn)
-          if isinstance(safe_options, int):
-            game["safe_options_sum"] += safe_options
-            game["safe_options_count"] += 1
-            if safe_options <= 2:
-              game["pressure_turns"] += 1
-
-          day_games[day].add(game_id)
+        day_games[day].add(game_id)
 
     game_summaries = []
     for game in games.values():
@@ -225,7 +186,7 @@ if __name__ == "__main__":
     "--input",
     action="append",
     required=True,
-    help="Input JSONL file, directory, or glob pattern. Repeat for multiple inputs.",
+    help="Input JSONL/JSONL.GZ file, directory, or glob pattern. Repeat for multiple inputs.",
   )
   parser.add_argument(
     "--output",