Skip to content

Commit

Permalink
fix: Special case getsizeof() for PyPy
Browse files Browse the repository at this point in the history
  • Loading branch information
akaihola committed Sep 18, 2024
1 parent 4a555ae commit e2d5f77
Showing 1 changed file with 12 additions and 6 deletions.
18 changes: 12 additions & 6 deletions pgtricks/mergesort.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""

Expand All @@ -25,18 +31,18 @@ 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:
"""Append a line to the set of lines to be sorted."""
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()

Expand All @@ -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."""
Expand Down

0 comments on commit e2d5f77

Please sign in to comment.