⚡️ Speed up mget() by 68% in libs/langchain/langchain/storage/file_system.py

[codeflash-ai]

📄 mget() in libs/langchain/langchain/storage/file_system.py

📈 Performance went up by 68% (0.68x faster)

⏱️ Runtime went down from 14075.68μs to 8401.84μs

Explanation and details

(click to show)

Your already well optimized Python code is ensuring file paths are valid, and checking for their existence. However, there are a few improvements we can implement while not altering function names/sigs and preserving output exactly as before.

First, pre-compile the regular expression outside of your get_full_path function. The re.compile() function reduces overhead since the pattern string isn’t read and converted to a regular expression object for every iteration, as opposed to re.match(r"^[a-zA-Z0-9.-/]+$", key).

Second, optimise the mget function to reduce the number of disk accesses. The os.path.exists(key) function requires a disk access, as well as os.path.read_bytes(). We could reduce this by utilising a try/except block in the value assignment, catching the FileNotFoundError exception.

Here's the optimized code.

This piece of code will have a faster runtime by avoiding redundant disk accesses and regex compilations.

Correctness verification

The new optimized code was tested for correctness. The results are listed below.

✅ 0 Passed − ⚙️ Existing Unit Tests

✅ 0 Passed − 🎨 Inspired Regression Tests

✅ 7 Passed − 🌀 Generated Regression Tests

(click to show generated tests)

# imports
import pytest
from pathlib import Path
import os
from typing import Sequence, List, Optional, Union

# Assuming the InvalidKeyException is defined somewhere in the code
class InvalidKeyException(Exception):
    pass

# Assuming the ByteStore is a base class defined somewhere in the code
class ByteStore:
    pass

# function to test
# ... (The provided _get_full_path and LocalFileStore classes go here) ...

# unit tests
class TestLocalFileStore:
    @pytest.fixture
    def file_store(self, tmp_path):
        """Fixture to create a LocalFileStore with a temporary directory."""
        store = LocalFileStore(root_path=tmp_path)
        # Create some files for testing
        (tmp_path / "file1.txt").write_text("content1")
        (tmp_path / "subdir").mkdir()
        (tmp_path / "subdir/file2.bin").write_bytes(b"content2")
        return store

    def test_valid_keys_existing_files(self, file_store, tmp_path):
        """Test valid keys with existing files."""
        keys = ["file1.txt", "subdir/file2.bin"]
        expected = [b"content1", b"content2"]
        assert file_store.mget(keys) == expected

    def test_valid_keys_non_existing_files(self, file_store):
        """Test valid keys with non-existing files."""
        keys = ["nonexistent.txt"]
        assert file_store.mget(keys) == [None]

    def test_mixed_existing_and_non_existing_files(self, file_store, tmp_path):
        """Test a mix of existing and non-existing files."""
        keys = ["file1.txt", "does_not_exist.bin"]
        expected = [b"content1", None]
        assert file_store.mget(keys) == expected

    def test_invalid_keys(self, file_store):
        """Test keys with invalid characters."""
        keys = ["invalid/key*name"]
        with pytest.raises(InvalidKeyException):
            file_store.mget(keys)

    def test_empty_keys_sequence(self, file_store):
        """Test an empty sequence of keys."""
        assert file_store.mget([]) == []

    def test_keys_with_dots(self, file_store, tmp_path):
        """Test keys with '.' and '..'."""
        (tmp_path / "file3.txt").write_text("content3")
        keys = ["./file1.txt", "subdir/../file3.txt"]
        expected = [b"content1", b"content3"]
        assert file_store.mget(keys) == expected

    def test_keys_with_traversal_attempt(self, file_store):
        """Test keys that attempt directory traversal."""
        keys = ["../file1.txt"]
        with pytest.raises(InvalidKeyException):
            file_store.mget(keys)

    def test_keys_with_leading_trailing_slashes(self, file_store, tmp_path):
        """Test keys with leading or trailing slashes."""
        (tmp_path / "file4.txt").write_text("content4")
        keys = ["/file1.txt", "subdir/file2.bin/"]
        expected = [b"content1", b"content2"]
        assert file_store.mget(keys) == expected

    def test_large_number_of_keys(self, file_store, tmp_path):
        """Test a large number of keys."""
        # Create 100 files
        for i in range(100):
            (tmp_path / f"file{i}.txt").write_text(f"content{i}")
        keys = [f"file{i}.txt" for i in range(100)]
        expected = [f"content{i}".encode() for i in range(100)]
        assert file_store.mget(keys) == expected

    def test_boundary_path_length(self, file_store, tmp_path):
        """Test keys at the boundary of the filesystem path length limit."""
        # This test is system-dependent; some systems may not have a path length limit
        long_filename = "a" * 255
        (tmp_path / long_filename).write_text("long content")
        keys = [long_filename]
        expected = [b"long content"]
        assert file_store.mget(keys) == expected

[aphexcx]

This is a code replacer bug! Retest with current version (this ran with 0.3.4)

It tacked the updated (dependent) function at the end of the file instead of replacing it in the class.