refactor: extract shared frame upload helpers (DRY)

[PythonFZ]

Extract common logic from read_file and _on_filesystem_load into shared helper functions:

• upload_frames_to_room(): batch processing with lock and progress
• apply_adaptive_resolution(): adaptive sphere resolution for large systems

This removes ~60 lines of duplicated code and fixes a bug in _on_filesystem_load where it tried to modify a frozen Pydantic model.

Summary by CodeRabbit

• Performance & Improvements
  • Optimized frame loading with batch processing and enhanced progress tracking
  • Improved rendering performance through adaptive resolution adjustments based on particle count
  • Streamlined file loading with better error handling

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

Two helper functions for frame loading and adaptive resolution were introduced, consolidating frame upload logic and geometry resolution adjustments into shared utilities. The filesystem load flow refactored to delegate frame streaming and resolution adaptation to these helpers instead of inlined implementations.

Changes

Cohort / File(s)	Summary
Frame Loading Helpers src/zndraw/app/tasks.py	Added upload_frames_to_room() to stream frames in batches with progress tracking (capped at 99%). Added apply_adaptive_resolution() to compute and apply adaptive sphere geometry resolution based on max particle count. Refactored read_file() to use these helpers, removing inlined adaptive-resolution logic.
Socket Manager Integration src/zndraw/socket_manager.py	Refactored _on_filesystem_load() to remove in-flight progress tracking and use shared helpers. Introduced dynamic backend selection via FORMAT_BACKENDS based on file extension. Replaced monolithic per-format pipeline with dedicated branches for ZnH5MD, ASE-DB, and generic formats. Now constructs frame iterators and delegates uploading/resolution to upload_frames_to_room() and apply_adaptive_resolution(). Simplified error handling and removed nested progress task contexts.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 Frames now flow in batches neat,
Helper hops make code complete,
Adaptive spheres spin round with grace,
Progress marches to 99's place! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 57.14% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main change: extracting shared frame upload helpers to reduce code duplication (DRY principle).

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

🧹 Recent nitpick comments

src/zndraw/app/tasks.py (2)

70-76: Type hints missing for vis and frame_iterator parameters.

Per coding guidelines, type hints should be used wherever possible. The vis and frame_iterator parameters lack type annotations.

♻️ Suggested improvement

+import typing as t
+
+if t.TYPE_CHECKING:
+    from zndraw.zndraw import ZnDraw
+    from ase import Atoms
+
 def upload_frames_to_room(
-    vis,
-    frame_iterator,
+    vis: "ZnDraw",
+    frame_iterator: t.Iterator["Atoms"],
     batch_size: int = 10,
     total_expected_frames: int | None = None,
     progress_description: str = "Loading frames",
 ) -> tuple[int, int]:

---

115-149: Type hint missing for vis parameter; broad exception handling is acceptable here.

The vis parameter should have a type hint per coding guidelines. The broad except Exception (flagged by static analysis) is acceptable in this context since adaptive resolution is a non-critical optimization—failing gracefully with a warning is the right behavior.

♻️ Add type hint

+if t.TYPE_CHECKING:
+    from zndraw.zndraw import ZnDraw
+
-def apply_adaptive_resolution(vis, max_particles: int) -> None:
+def apply_adaptive_resolution(vis: "ZnDraw", max_particles: int) -> None:

src/zndraw/socket_manager.py (3)

570-583: Consider explicit cleanup of h5py.File for robustness.

The h5py.File object at line 573 is not explicitly closed. While it will be garbage collected after the iterator is consumed and the with fs.open() block exits, explicitly closing it ensures resources are released promptly—especially important if an exception occurs during iteration.

♻️ Suggested improvement

                 if backend_name == "ZnH5MD":
                     if step is not None and step <= 0:
                         raise ValueError("Step must be a positive integer")
-                    h5_file = h5py.File(f, "r")
-                    io = znh5md.IO(file_handle=h5_file)
-                    n_frames = len(io)
-
-                    if start is not None and start >= n_frames:
-                        raise ValueError(f"Start frame {start} exceeds file length")
-
-                    s = slice(start, stop, step)
-                    _start, _stop, _step = s.indices(n_frames)
-                    total_expected_frames = (_stop - _start + _step - 1) // _step
-                    frame_iterator = itertools.islice(io, _start, _stop, _step)
+                    h5_file = h5py.File(f, "r")
+                    try:
+                        io = znh5md.IO(file_handle=h5_file)
+                        n_frames = len(io)
+
+                        if start is not None and start >= n_frames:
+                            raise ValueError(f"Start frame {start} exceeds file length")
+
+                        s = slice(start, stop, step)
+                        _start, _stop, _step = s.indices(n_frames)
+                        total_expected_frames = (_stop - _start + _step - 1) // _step
+                        # Materialize frames before closing h5_file
+                        frame_iterator = list(itertools.islice(io, _start, _stop, _step))
+                    finally:
+                        h5_file.close()

Note: This materializes frames into a list, which increases memory usage but ensures proper cleanup. Alternatively, you could keep the current lazy approach if memory is a concern, accepting that cleanup relies on GC.

---

605-612: Inconsistent behavior: missing rows are silently skipped here but logged in tasks.py.

In tasks.py (lines 282-291), when a database row is not found, a warning is logged via vis.log(). Here, the KeyError is silently caught and the row is skipped without notification. Consider adding consistent logging for better observability.

♻️ Add warning log for consistency

                     def db_iterator():
                         for row_id in selected_ids:
                             try:
                                 yield db.get(id=row_id).toatoms()
                             except KeyError:
+                                log.warning(f"Row ID {row_id} not found in database, skipping")
                                 continue

---

543-544: Inconsistent error message detail compared to other filesystem handlers.

This error message is simplified compared to _on_filesystem_list (lines 432-436) and _on_filesystem_metadata (lines 681-685), which include available keys and the public flag for debugging. Consider maintaining consistency.

♻️ Restore detailed error message

             if fs_key not in self.zndraw._filesystems:
-                raise ValueError(f"Filesystem '{fs_name}' not found")
+                available_keys = list(self.zndraw._filesystems.keys())
+                raise ValueError(
+                    f"Filesystem '{fs_name}' (key: '{fs_key}') not found. "
+                    f"Available: {available_keys}. "
+                    f"Public flag: {public}, Room: {self.zndraw.room}"
+                )

---

📜 Recent review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 062ff05 and 5025ed3.

📒 Files selected for processing (2)

• src/zndraw/app/tasks.py
• src/zndraw/socket_manager.py

🧰 Additional context used

📓 Path-based instructions (1)

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: You MUST NEVER @pytest.mark.xfail or similar - all tests must pass!
All default values (e.g., camera, particles) must be defined exclusively within the Pydantic model. Do not scatter fallback logic throughout the codebase.
Do not perform null checks combined with hardcoded literals for default values. Rely entirely on the schema to populate default values during initialization.
You can not use LUA scripts with Redis!
If sensible, implement collections.abc interfaces for your classes, such as MutableMapping or MutableSequence.
Use numpy style docstrings in Python code.
Docstrings must be concise and to the point.
Use type hints wherever possible in Python. Use list[int|float] | None instead of t.Optional[t.List[int|float]]!
Imports should always be at the top of the file in Python.

Files:

• src/zndraw/app/tasks.py
• src/zndraw/socket_manager.py

🧬 Code graph analysis (2)

src/zndraw/app/tasks.py (4)

src/zndraw/zndraw.py (12)

• progress_tracker (1734-1755)
• atoms (757-759)
• atoms (762-764)
• update (346-360)
• geometries (573-574)
• step (722-724)
• step (727-740)
• get (213-226)
• get (926-926)
• get (928-930)
• get (932-934)
• get (963-1056)

src/zndraw/geometries/sphere.py (1)

• Sphere (13-62)

src/zndraw/storage/asebytes_backend.py (1)

• get (46-103)

src/zndraw/frame_cache.py (1)

• get (16-17)

src/zndraw/socket_manager.py (1)

src/zndraw/app/tasks.py (2)

• apply_adaptive_resolution (115-149)
• upload_frames_to_room (70-112)

🪛 Ruff (0.14.11)

src/zndraw/app/tasks.py

148-148: Do not catch blind exception: Exception

(BLE001)

---

228-228: Abstract raise to an inner function

(TRY301)

---

228-228: Avoid specifying long messages outside the exception class

(TRY003)

---

236-236: Abstract raise to an inner function

(TRY301)

---

236-236: Avoid specifying long messages outside the exception class

(TRY003)

---

260-260: Abstract raise to an inner function

(TRY301)

---

260-260: Avoid specifying long messages outside the exception class

(TRY003)

---

270-270: Abstract raise to an inner function

(TRY301)

---

270-270: Avoid specifying long messages outside the exception class

(TRY003)

src/zndraw/socket_manager.py

544-544: Abstract raise to an inner function

(TRY301)

---

544-544: Avoid specifying long messages outside the exception class

(TRY003)

---

572-572: Abstract raise to an inner function

(TRY301)

---

572-572: Avoid specifying long messages outside the exception class

(TRY003)

---

578-578: Abstract raise to an inner function

(TRY301)

---

578-578: Avoid specifying long messages outside the exception class

(TRY003)

---

591-591: Abstract raise to an inner function

(TRY301)

---

591-591: Avoid specifying long messages outside the exception class

(TRY003)

---

593-593: Abstract raise to an inner function

(TRY301)

---

593-593: Avoid specifying long messages outside the exception class

(TRY003)

---

600-600: Abstract raise to an inner function

(TRY301)

---

600-600: Avoid specifying long messages outside the exception class

(TRY003)

---

642-646: Consider moving this statement to an else block

(TRY300)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)

• GitHub Check: pytest (3.12, ubuntu-latest)
• GitHub Check: pytest (3.13, ubuntu-latest)
• GitHub Check: pytest (3.11, ubuntu-latest)

🔇 Additional comments (4)

src/zndraw/app/tasks.py (2)

303-314: LGTM!

The refactor cleanly delegates frame uploading and progress tracking to the shared helper. The conditional check for frame_iterator truthiness handles edge cases properly.

---

327-329: LGTM!

Correctly guards adaptive resolution behind a frame count check—no point adjusting rendering for empty uploads.

src/zndraw/socket_manager.py (2)

524-525: LGTM!

Good consolidation—importing the shared helpers from tasks.py eliminates code duplication and ensures consistent behavior between the Celery task and socket-based loading.

---

623-640: LGTM!

The frame upload and adaptive resolution logic correctly delegates to the shared helpers, maintaining consistency with the Celery task implementation. The success response includes the frame count for proper feedback.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov-commenter]

Codecov Report

❌ Patch coverage is 68.21192% with 48 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.02%. Comparing base (062ff05) to head (5025ed3).
⚠️ Report is 47 commits behind head on main.

Files with missing lines	Patch %	Lines
src/zndraw/socket_manager.py	41.37%	34 Missing ⚠️
src/zndraw/app/tasks.py	84.94%	14 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #854      +/-   ##
==========================================
+ Coverage   79.97%   80.02%   +0.04%     
==========================================
  Files         165      165              
  Lines       20143    20128      -15     
==========================================
- Hits        16110    16107       -3     
+ Misses       4033     4021      -12     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.