Merge pull request #29 from bemxio/dumping

dumping features for all gamemodes

README.md

@@ -37,6 +37,25 @@ replay = parse_replay(lzma_string, pure_lzma=True)
 
 Note that if you use the `/get_replay` endpoint to retrieve a replay, you must decode the response before passing it to osrparse, as the response is encoded in base 64 by default.
 
+### Dumping
+
+To dump a replay to a filepath:
+```python
+from osrparse import dump_replay_file
+
+# ...some parsing code here
+dump_replay_file(replay, "path/to/osr.osr")
+```
+
+To dump a replay into a variable:
+
+```python
+from osrparse import dump_replay
+
+# ...some parsing code here
+osr_content = dump_replay(replay)
+```
+
 ### Attributes
 
 `Replay` objects have the following attibutes:

osrparse/__init__.py

@@ -6,6 +6,7 @@
 __version__ = "5.0.0"
 
 
-__all__ = ["GameMode", "Mod", "parse_replay_file", "parse_replay", "Replay",
-    "ReplayEvent", "Key", "ReplayEventOsu", "ReplayEventTaiko",
-    "ReplayEventMania", "ReplayEventCatch", "KeyTaiko", "KeyMania"]
+__all__ = ["GameMode", "Mod", "parse_replay_file", "parse_replay", 
+    "Replay", "ReplayEvent", "Key",
+    "ReplayEventOsu", "ReplayEventTaiko", "ReplayEventMania",
+    "ReplayEventCatch", "KeyTaiko", "KeyMania"]

osrparse/dump.py

@@ -0,0 +1,94 @@
+import lzma
+import struct
+
+from osrparse.utils import (ReplayEventOsu, ReplayEventTaiko, ReplayEventCatch,
+    ReplayEventMania)
+
+def pack_byte(data: int):
+    return struct.pack("<B", data)
+
+def pack_short(data: int):
+    return struct.pack("<H", data)
+
+def pack_int(data: int):
+    return struct.pack("<I", data)
+
+def pack_long(data: int):
+    return struct.pack("<Q", data)
+
+def pack_ULEB128(data):
+    # taken from https://github.com/mohanson/leb128
+    r, i = [], len(data)
+
+    while True:
+        byte = i & 0x7f
+        i = i >> 7
+
+        if (i == 0 and byte & 0x40 == 0) or (i == -1 and byte & 0x40 != 0):
+            r.append(byte)
+            return b"".join(map(pack_byte, r))
+
+        r.append(0x80 | byte)
+
+def pack_string(data: str):
+    if data:
+        return pack_byte(11) + pack_ULEB128(data) + data.encode("utf-8")
+    return pack_byte(11) + pack_byte(0)
+
+def dump_timestamp(replay):
+    # windows ticks starts at year 0001, in contrast to unix time (1970).
+    # 62135596800 is the number of seconds between these two years and is added
+    # to account for this difference.
+    # The factor of 10000000 converts seconds to ticks.
+    ticks = (62135596800 + replay.timestamp.timestamp()) * 10000000
+    ticks = int(ticks)
+    return pack_long(ticks)
+
+
+def dump_replay_data(replay):
+    replay_data = ""
+    for event in replay.play_data:
+        if isinstance(event, ReplayEventOsu):
+            replay_data += f"{event.time_delta}|{event.x}|{event.y}|{event.keys.value},"
+        elif isinstance(event, ReplayEventTaiko):
+            replay_data += f"{event.time_delta}|{event.x}|0|{event.keys.value},"
+        elif isinstance(event, ReplayEventCatch):
+            replay_data += f"{event.time_delta}|{event.x}|0|{int(event.dashing)},"
+        elif isinstance(event, ReplayEventMania):
+            replay_data += f"{event.time_delta}|{event.keys}|0|0,"
+
+    filters = [{"id": lzma.FILTER_LZMA1, "dict_size": 1 << 21, "mode": lzma.MODE_FAST}]
+    compressed = lzma.compress(replay_data.encode("ascii"), format=lzma.FORMAT_ALONE, filters=filters)
+
+    return pack_int(len(compressed)) + compressed
+
+
+def dump_replay(replay):
+    data = b""
+
+    data += pack_byte(replay.game_mode.value)
+    data += pack_int(replay.game_version)
+    data += pack_string(replay.beatmap_hash)
+
+    data += pack_string(replay.player_name)
+    data += pack_string(replay.replay_hash)
+
+    data += pack_short(replay.number_300s)
+    data += pack_short(replay.number_100s)
+    data += pack_short(replay.number_50s)
+    data += pack_short(replay.gekis)
+    data += pack_short(replay.katus)
+    data += pack_short(replay.misses)
+
+    data += pack_int(replay.score)
+    data += pack_short(replay.max_combo)
+    data += pack_byte(replay.is_perfect_combo)
+
+    data += pack_int(replay.mod_combination.value)
+    data += pack_string(replay.life_bar_graph)
+    data += dump_timestamp(replay)
+
+    data += dump_replay_data(replay)
+    data += pack_long(replay.replay_id)
+
+    return data

osrparse/replay.py

@@ -1,10 +1,12 @@
 import lzma
 import struct
-import datetime
+from datetime import datetime, timezone, timedelta
 from typing import List
+from io import TextIOWrapper
 
 from osrparse.utils import (Mod, GameMode, ReplayEvent, ReplayEventOsu,
     ReplayEventCatch, ReplayEventMania, ReplayEventTaiko)
+from osrparse.dump import dump_replay
 
 class Replay:
     # first version with rng seed value added as the last frame in the lzma data
@@ -112,7 +114,8 @@ def _parse_life_bar_graph(self, replay_data):
     def _parse_timestamp_and_replay_length(self, replay_data):
         format_specifier = "<qi"
         (t, self.replay_length) = struct.unpack_from(format_specifier, replay_data, self.offset)
-        self.timestamp = datetime.datetime.min + datetime.timedelta(microseconds=t/10)
+        self.timestamp = datetime.min + timedelta(microseconds=t/10)
+        self.timestamp = self.timestamp.replace(tzinfo=timezone.utc)
         self.offset += struct.calcsize(format_specifier)
 
     def _parse_play_data(self, replay_data):
@@ -176,3 +179,16 @@ def _parse_replay_id(self, replay_data):
             format_specifier = "<l"
             replay_id = struct.unpack_from(format_specifier, replay_data, self.offset)
         self.replay_id = replay_id[0]
+
+    def dump(self, file=None):
+        dumped = dump_replay(self)
+
+        if not file:
+            return dumped
+
+        # allow either a live file object, or a path to a file
+        if isinstance(file, TextIOWrapper):
+            file.write(dumped)
+        else:
+            with open(file, "wb") as f:
+                f.write(dumped)

tests/test_dumping.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+from unittest import TestCase
+from tempfile import TemporaryDirectory
+
+from osrparse import parse_replay_file
+
+
+RES = Path(__file__).parent / "resources"
+
+
+class TestDumping(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.replay = parse_replay_file(RES / "replay.osr")
+
+    def test_dumping(self):
+        with TemporaryDirectory() as tempdir:
+            r2_path = Path(tempdir) / "dumped.osr"
+            self.replay.dump(r2_path)
+            r2 = parse_replay_file(r2_path)
+
+        # `replay_length` is intentionally not tested for equality here, as the
+        # length of the compressed replay data may change after dumping due to
+        # varying lzma settings.
+        attrs = ["game_mode", "game_version", "beatmap_hash", "player_name",
+            "replay_hash", "number_300s", "number_100s", "number_50s", "gekis",
+            "katus", "misses", "score", "max_combo", "is_perfect_combo",
+            "mod_combination", "life_bar_graph", "timestamp", "play_data",
+            "replay_id"]
+        for attr in attrs:
+            self.assertEqual(getattr(self.replay, attr), getattr(r2, attr),
+                f"{attr} is wrong")

tests/test_replay.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 from unittest import TestCase
-import datetime
+from datetime import datetime, timezone
 from osrparse import (parse_replay, parse_replay_file, ReplayEventOsu, GameMode,
     Mod, ReplayEventTaiko, ReplayEventCatch, ReplayEventMania)
 
@@ -60,7 +60,7 @@ def test_mod_combination(self):
 
     def test_timestamp(self):
         for replay in self._replays:
-            self.assertEqual(replay.timestamp, datetime.datetime(2013, 2, 1, 16, 31, 34), "Timestamp is wrong")
+            self.assertEqual(replay.timestamp, datetime(2013, 2, 1, 16, 31, 34, tzinfo=timezone.utc), "Timestamp is wrong")
 
     def test_play_data(self):
         for replay in self._replays: