diff --git a/README.md b/README.md index d1fb846..b4ff416 100644 --- a/README.md +++ b/README.md @@ -969,4 +969,4 @@ pull request to ReadStat first. [alchemyst](https://github.com/alchemyst): improvements to docstrings -[bmwiedemann](https://github.com/bmwiedemann), [toddrme2178 ](https://github.com/toddrme2178), [Martin Thorsen Ranang](https://github.com/mtr): improvements to source code +[bmwiedemann](https://github.com/bmwiedemann), [toddrme2178 ](https://github.com/toddrme2178), [Martin Thorsen Ranang](https://github.com/mtr), [Eirik Stavestrand](https://github.com/eirki): improvements to source code diff --git a/pyreadstat/_readstat_parser.pyx b/pyreadstat/_readstat_parser.pyx index 802fd4c..3de0814 100644 --- a/pyreadstat/_readstat_parser.pyx +++ b/pyreadstat/_readstat_parser.pyx @@ -852,7 +852,6 @@ cdef int handle_open(const char *u8_path, void *io_ctx) except READSTAT_HANDLER_ return -1 -cdef object _file_object_ctx = None cdef int pyobject_open_handler(const char *path, void *io_ctx) noexcept: """File is already open - this is a no-op""" @@ -864,13 +863,12 @@ cdef int pyobject_close_handler(void *io_ctx) noexcept: cdef ssize_t pyobject_read_handler(void *buf, size_t nbyte, void *io_ctx) noexcept: """Bridge Python file.read() to C read operation""" - global _file_object_ctx cdef bytes data cdef ssize_t bytes_read cdef char *data_ptr - + cdef object file_obj = io_ctx + try: - file_obj = _file_object_ctx data = file_obj.read(nbyte) bytes_read = len(data) if bytes_read > 0: @@ -882,18 +880,17 @@ cdef ssize_t pyobject_read_handler(void *buf, size_t nbyte, void *io_ctx) noexce cdef readstat_off_t pyobject_seek_handler(readstat_off_t offset, readstat_io_flags_t whence, void *io_ctx) noexcept: """Bridge Python file.seek() to C seek operation""" - global _file_object_ctx cdef int py_whence - + cdef object file_obj = io_ctx + try: - file_obj = _file_object_ctx if whence == READSTAT_SEEK_SET: py_whence = 0 elif whence == READSTAT_SEEK_CUR: py_whence = 1 else: # READSTAT_SEEK_END py_whence = 2 - + file_obj.seek(offset, py_whence) return file_obj.tell() except: @@ -916,11 +913,9 @@ cdef void check_exit_status(readstat_error_t retcode) except *: cdef void run_readstat_parser(char * filename, data_container data, py_file_extension file_extension, long row_limit, long row_offset, object file_obj=None) except *: """ Runs the parsing of the file by readstat library. - + If file_obj is provided, it will be used instead of filename for I/O operations. """ - global _file_object_ctx - cdef readstat_parser_t *parser cdef readstat_error_t error cdef readstat_metadata_handler metadata_handler @@ -959,11 +954,12 @@ cdef void run_readstat_parser(char * filename, data_container data, py_file_exte # Set up custom I/O handlers for file objects if file_obj is not None: - _file_object_ctx = file_obj + io_ctx = file_obj open_handler = pyobject_open_handler close_handler = pyobject_close_handler read_handler = pyobject_read_handler seek_handler = pyobject_seek_handler + readstat_set_io_ctx(parser, io_ctx) readstat_set_open_handler(parser, open_handler) readstat_set_close_handler(parser, close_handler) readstat_set_read_handler(parser, read_handler) diff --git a/tests/test_narwhalified.py b/tests/test_narwhalified.py index 3d9ec46..cc6f8df 100644 --- a/tests/test_narwhalified.py +++ b/tests/test_narwhalified.py @@ -16,12 +16,14 @@ # ############################################################################# from datetime import datetime, timedelta, date +from concurrent.futures import ThreadPoolExecutor import unittest import os import sys import shutil import multiprocessing as mp import tempfile +import time import zipfile import io @@ -1378,6 +1380,32 @@ def test_read_sav_file_handle(self): self.assertEqual(len(df.columns), len(self.df_pandas.columns)) self.assertEqual(len(df), len(self.df_pandas)) + def test_read_sav_bytesio_threads(self): + """Test reading SAV file from file-like object in multiple threads at once (tests thread safety)""" + + class SlowBytesIO(io.BytesIO): + """A BytesIO that sleeps a bit on each read. This subclass is necessary because we want + to test paralell threads reading at the same time, but the test data is so small that + the reads are too fast to overlap. """ + def read(self, *args, **kwargs): + time.sleep(0.001) + return super().read(*args, **kwargs) + + def read_sav_file(buffer): + df, meta = pyreadstat.read_sav(buffer, output_format=self.backend) + return df, meta + + sav_file = os.path.join(self.basic_data_folder, "sample.sav") + with open(sav_file, "rb") as f: + file_bytes = f.read() + num_threads = 5 + buffers = [SlowBytesIO(file_bytes) for _ in range(num_threads)] + with ThreadPoolExecutor(max_workers=num_threads) as executor: + results = list(executor.map(read_sav_file, buffers)) + for df, meta in results: + self.assertEqual(len(df.columns), len(self.df_pandas.columns)) + self.assertEqual(len(df), len(self.df_pandas)) + self.assertListEqual(list(df.columns), list(self.df_pandas.columns)) def test_read_sav_bytesio(self): """Test reading SAV file from BytesIO (simulates remote/streaming data)"""