From d2acfe46ca54777b50f862a6b62f5126c1d4544e Mon Sep 17 00:00:00 2001
From: Matt Clarke <mjclarke94@gmail.com>
Date: Fri, 30 Apr 2021 23:21:05 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=91=8C=20IMPROVE:=20Write=20subset=20of?=
 =?UTF-8?q?=20trajectory=20frames=20to=20disk.=20(#22)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Chris Sewell <chrisj_sewell@hotmail.com>
---
 .github/workflows/tests.yml           |  2 +-
 aiida_lammps/data/trajectory.py       | 23 ++++++++++++++++-------
 aiida_lammps/tests/test_trajectory.py |  7 +++++++
 codecov.yml                           | 10 ++++++++++
 4 files changed, 34 insertions(+), 8 deletions(-)
 create mode 100644 codecov.yml

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 1c726ea..4166e84 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -90,7 +90,7 @@ jobs:
         pytest  --lammps-exec lmp_serial --cov=aiida_lammps --cov-report=xml --cov-report=term-missing
 
     - name: Upload to Codecov
-      if: matrix.os == 'ubuntu-latest' && matrix.python-version == 3.7
+      if: matrix.python-version == 3.8
       uses: codecov/codecov-action@v1
       with:
         name: pytests-lammps
diff --git a/aiida_lammps/data/trajectory.py b/aiida_lammps/data/trajectory.py
index 468957b..6b8d562 100644
--- a/aiida_lammps/data/trajectory.py
+++ b/aiida_lammps/data/trajectory.py
@@ -169,13 +169,19 @@ def get_step_data(self, step_idx):
         """Return parsed data, for a specific trajectory step."""
         return parse_step(self.get_step_string(step_idx).splitlines())
 
-    def iter_step_strings(self):
-        """Yield the content string, for a each trajectory step."""
+    def iter_step_strings(self, steps=None):
+        """Yield the content string, for each trajectory step."""
+
+        if steps is None:
+            steps = range(self.number_steps)
+        elif isinstance(steps, int):
+            steps = range(0, self.number_steps, steps)
+
         with self.open(self.get_attribute("traj_filename"), mode="rb") as handle:
             with ZipFile(
                 handle, "r", self.get_attribute("compression_method")
             ) as zip_file:
-                for step_idx in range(self.number_steps):
+                for step_idx in steps:
                     zip_name = "{}{}".format(self.get_attribute("zip_prefix"), step_idx)
                     with zip_file.open(zip_name) as step_file:
                         content = step_file.read()
@@ -213,8 +219,11 @@ def get_step_structure(
             original_structure=original_structure,
         )
 
-    def write_as_lammps(self, handle):
-        """Write out the lammps trajectory to file."""
-        for string in self.iter_step_strings():
+    def write_as_lammps(self, handle, steps=None):
+        """Write out the lammps trajectory to file.
+        :param handle: a file handle, opened in "wb" mode
+        :param steps: a list of steps to write (default to all)
+        """
+        for string in self.iter_step_strings(steps=steps):
             handle.write(string)
-            handle.write("\n")
+            handle.write(b"\n")
diff --git a/aiida_lammps/tests/test_trajectory.py b/aiida_lammps/tests/test_trajectory.py
index 82d7f60..394c079 100644
--- a/aiida_lammps/tests/test_trajectory.py
+++ b/aiida_lammps/tests/test_trajectory.py
@@ -73,3 +73,10 @@ def test_lammpstraj_timesteps(db_test_app):
         900,
         1000,
     ]
+
+
+def test_write_as_lammps(db_test_app, tmp_path):
+    path = os.path.join(TEST_DIR, "input_files", "trajectory.lammpstrj")
+    data = LammpsTrajectory(path)
+    with tmp_path.joinpath("trajectory.lammpstrj").open(mode="wb") as handle:
+        data.write_as_lammps(handle)
diff --git a/codecov.yml b/codecov.yml
new file mode 100644
index 0000000..bbba85b
--- /dev/null
+++ b/codecov.yml
@@ -0,0 +1,10 @@
+coverage:
+  status:
+    project:
+      default:
+        target: 75%
+        threshold: 0.2%
+    patch:
+      default:
+        target: 70%
+        threshold: 0.2%