From 3af6c47691aadd4432d719ccdda81eacdefbbdc2 Mon Sep 17 00:00:00 2001
From: BC44 <63179336+BC44@users.noreply.github.com>
Date: Sun, 30 Aug 2020 22:59:37 -0400
Subject: [PATCH] Improve piece sizes

Improve piece size to avoid overly large torrent file size
---
 torrentool/torrent.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/torrentool/torrent.py b/torrentool/torrent.py
index c4f876e..6690e29 100644
--- a/torrentool/torrent.py
+++ b/torrentool/torrent.py
@@ -15,9 +15,22 @@
 
 _ITERABLE_TYPES = (list, tuple, set)
 
-
 TorrentFile = namedtuple('TorrentFile', ['name', 'length'])
 
+# path size in MiB (upper bound) : piece length
+PIECE_LENGTHS = {
+    50: 32,  # path size [0 - 50) MiB --> 32 kiB piece size
+    150: 64,  # path size [50 - 150) MiB --> 64 kiB piece size
+    350: 128,  # 128 kiB piece size
+    500: 256,  # 256 kiB piece size
+    1000: 512,  # 512 kiB piece size
+    2000: 1024,  # 1024 kiB piece size
+    4000: 2048,  # 2048 kiB piece size
+    8000: 4096,  # 4096 kiB piece size
+    16000: 8192,  # 8192 kiB piece size
+    16001: 16384  # 16384 kiB piece size
+}
+
 
 class Torrent:
     """Represents a torrent file, and exposes utilities to work with it."""
@@ -298,6 +311,13 @@ def to_string(self) -> bytes:
         """Returns bytes representing torrent file."""
         return Bencode.encode(self._struct)
 
+    def get_piece_size(self, size_data):
+        size_data_mb = size_data * 1024 * 1024
+        for size_mb in PIECE_LENGTHS:
+            if size_data_mb < size_mb:
+                return PIECE_LENGTHS[size_mb] * 1024
+        return PIECE_LENGTHS[16001] * 1024
+
     @classmethod
     def _get_target_files_info(cls, src_path: Path) -> Tuple[List[Tuple[str, int, List[str]]], int]:
         is_dir = src_path.is_dir()
@@ -346,12 +366,7 @@ def create_from(cls, src_path: Union[str, Path]) -> 'Torrent':
         # chunks_min = 1000
         # chunks_max = 2200
 
-        size_piece = size_min
-        if size_data > size_min:
-            size_piece = size_default
-
-        if size_piece > size_max:
-            size_piece = size_max
+        size_piece = self.get_piece_size(size_data)
 
         def read(filepath):
             with open(filepath, 'rb') as f: