From e2d5f77aa8ab96b0a1f9f17fee61aeec5ee105d5 Mon Sep 17 00:00:00 2001
From: Antti Kaihola <13725+akaihola@users.noreply.github.com>
Date: Thu, 25 Apr 2024 22:19:54 +0300
Subject: [PATCH] fix: Special case `getsizeof()` for PyPy

---
 pgtricks/mergesort.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/pgtricks/mergesort.py b/pgtricks/mergesort.py
index f22f020..5cccb52 100644
--- a/pgtricks/mergesort.py
+++ b/pgtricks/mergesort.py
@@ -2,12 +2,18 @@
 
 from __future__ import annotations
 
-import sys
+import platform
 from heapq import merge
 from tempfile import TemporaryFile
 from typing import IO, Any, Callable, Iterable, Iterator, cast
 
 
+if platform.python_implementation() == "PyPy":
+    getsizeof = len
+else:
+    from sys import getsizeof
+
+
 class MergeSort(Iterable[str]):
     """Merge sort implementation to handle large files by sorting them in partitions."""
 
@@ -25,7 +31,7 @@ def __init__(
         self._partitions: list[IO[bytes]] = []
         self._iterating: Iterable[str] | None = None
         self._buffer: list[str] = []
-        self._memory_counter: int = sys.getsizeof(self._buffer)
+        self._memory_counter: int = getsizeof(self._buffer)
         self._flush()
 
     def append(self, line: str) -> None:
@@ -33,10 +39,10 @@ def append(self, line: str) -> None:
         if self._iterating:
             message = "Can't append lines after starting to sort"
             raise ValueError(message)
-        self._memory_counter -= sys.getsizeof(self._buffer)
+        self._memory_counter -= getsizeof(self._buffer)
         self._buffer.append(line)
-        self._memory_counter += sys.getsizeof(self._buffer)
-        self._memory_counter += sys.getsizeof(line)
+        self._memory_counter += getsizeof(self._buffer)
+        self._memory_counter += getsizeof(line)
         if self._memory_counter >= self._max_memory:
             self._flush()
 
@@ -48,7 +54,7 @@ def _flush(self) -> None:
                 line.encode("UTF-8") for line in sorted(self._buffer, key=self._key)
             )
         self._buffer = []
-        self._memory_counter = sys.getsizeof(self._buffer)
+        self._memory_counter = getsizeof(self._buffer)
 
     def __next__(self) -> str:
         """Return the next line in the sorted list of lines."""