Merge pull request #157 from amgross/fix_repl_remove_redundant

Fix bug in REPL and aggregate duplicated code

Author: BrianPugh

src/littlefs/__main__.py

@@ -8,7 +8,7 @@
 from littlefs import LittleFS, __version__
 from littlefs.errors import LittleFSError
 from littlefs.repl import LittleFSRepl
-from littlefs.context import UserContextFile
+from littlefs.context import UserContextFile, UserContext
 
 # Dictionary mapping suffixes to their size in bytes
 _suffix_map = {
@@ -18,10 +18,12 @@
 }
 
 
-def _fs_from_args(args: argparse.Namespace, mount=True) -> LittleFS:
+def _fs_from_args(args: argparse.Namespace, block_count=None, mount=True, context: UserContext = None) -> LittleFS:
+    block_count = block_count if block_count is not None else getattr(args, "block_count", 0)
     return LittleFS(
+        context=context,
         block_size=args.block_size,
-        block_count=getattr(args, "block_count", 0),
+        block_count=block_count,
         name_max=args.name_max,
         mount=mount,
     )
@@ -105,11 +107,7 @@ def create(parser: argparse.ArgumentParser, args: argparse.Namespace) -> int:
     if args.compact:
         if args.verbose:
             print(f"Compacting... {fs.used_block_count} / {args.block_count}")
-        compact_fs = LittleFS(
-            block_size=args.block_size,
-            block_count=fs.used_block_count,
-            name_max=args.name_max,
-        )
+        compact_fs = _fs_from_args(args, block_count=fs.used_block_count)
         for root, dirs, files in fs.walk("/"):
             if not root.endswith("/"):
                 root += "/"
@@ -131,21 +129,37 @@ def create(parser: argparse.ArgumentParser, args: argparse.Namespace) -> int:
     return 0
 
 
-def _list(parser: argparse.ArgumentParser, args: argparse.Namespace) -> int:
-    """List LittleFS image contents."""
-    fs = _fs_from_args(args, mount=False)
-    fs.context.buffer = bytearray(args.source.read_bytes())
+def _mount_from_context(parser: argparse.ArgumentParser, args: argparse.Namespace, context: UserContext) -> LittleFS:
+    # Block count is 0 because we don't know the size of the real image yet, the source file may be compacted (with the create --compact option).
+    fs = _fs_from_args(args, block_count=0, mount=False, context=context)
     fs.mount()
 
     if args.verbose:
-        fs_size = len(fs.context.buffer)
+        input_image_size = context.in_size
+        actual_image_size = fs.block_count * args.block_size
         print("LittleFS Configuration:")
         print(f"  Block Size:  {args.block_size:9d}  /  0x{args.block_size:X}")
-        print(f"  Image Size:  {fs_size:9d}  /  0x{fs_size:X}")
+        if input_image_size != actual_image_size:
+            print(f"  Image Size:  {actual_image_size:9d}  /  0x{actual_image_size:X}")
+            print(f"  Input Image Size (compacted): {input_image_size:9d}  /  0x{input_image_size:X}")
+        else:
+            print(f"  Image Size:  {input_image_size:9d}  /  0x{input_image_size:X}")
         print(f"  Block Count: {fs.block_count:9d}")
         print(f"  Name Max:    {args.name_max:9d}")
         print(f"  Image:       {args.source}")
 
+    return fs
+
+
+def _list(parser: argparse.ArgumentParser, args: argparse.Namespace) -> int:
+    """List LittleFS image contents."""
+    source: Path = args.source
+    if not source.is_file():
+        parser.error(f"Source image '{source}' does not exist.")
+    context = UserContext(buffer=bytearray(source.read_bytes()))
+
+    fs = _mount_from_context(parser, args, context)
+
     for root, dirs, files in fs.walk("/"):
         if not root.endswith("/"):
             root += "/"
@@ -158,18 +172,12 @@ def _list(parser: argparse.ArgumentParser, args: argparse.Namespace) -> int:
 
 def extract(parser: argparse.ArgumentParser, args: argparse.Namespace) -> int:
     """Extract LittleFS image contents to a directory."""
-    fs = _fs_from_args(args, mount=False)
-    fs.context.buffer = bytearray(args.source.read_bytes())
-    fs.mount()
+    source: Path = args.source
+    if not source.is_file():
+        parser.error(f"Source image '{source}' does not exist.")
+    context = UserContext(buffer=bytearray(source.read_bytes()))
 
-    if args.verbose:
-        fs_size = len(fs.context.buffer)
-        print("LittleFS Configuration:")
-        print(f"  Block Size:  {args.block_size:9d}  /  0x{args.block_size:X}")
-        print(f"  Image Size:  {fs_size:9d}  /  0x{fs_size:X}")
-        print(f"  Block Count: {fs.block_count:9d}")
-        print(f"  Name Max:    {args.name_max:9d}")
-        print(f"  Image:       {args.source}")
+    fs = _mount_from_context(parser, args, context)
 
     root_dest = args.destination.absolute()
     if not root_dest.exists():
@@ -202,36 +210,19 @@ def extract(parser: argparse.ArgumentParser, args: argparse.Namespace) -> int:
 
 def repl(parser: argparse.ArgumentParser, args: argparse.Namespace) -> int:
     """Inspect an existing LittleFS image through an interactive shell."""
-
     source: Path = args.source
     if not source.is_file():
         parser.error(f"Source image '{source}' does not exist.")
+    context = UserContextFile(str(source))  # In repl we want context to be the file itself, so commands will change it
 
-    image_size = source.stat().st_size
-    if not image_size or image_size % args.block_size:
-        parser.error(
-            f"Image size ({image_size} bytes) is not a multiple of the supplied block size ({args.block_size})."
-        )
-
-    block_count = image_size // args.block_size
-    if block_count == 0:
-        parser.error("Image is smaller than a single block; cannot mount.")
-
-    context = UserContextFile(str(source))
-    fs = LittleFS(
-        context=context,
-        block_size=args.block_size,
-        block_count=block_count,
-        name_max=args.name_max,
-        mount=False,
-    )
-
-    shell = LittleFSRepl(fs)
     try:
         try:
-            shell.do_mount()
+            fs = _mount_from_context(parser, args, context)
         except LittleFSError as exc:
             parser.error(f"Failed to mount '{source}': {exc}")
+
+        shell = LittleFSRepl(fs)
+
         shell.cmdloop()
     finally:
         if shell._mounted:
@@ -367,6 +358,7 @@ def add_command(handler, name="", help=""):
 
     return parser
 
+
 # Getting argv optionally from the caller to enable call from python (generally for testing, but could be used for other purposes)
 def main(argv=None):
     if argv is None:
@@ -379,4 +371,3 @@ def main(argv=None):
 
 if __name__ == "__main__":
     sys.exit(main())
-

src/littlefs/context.py

@@ -10,8 +10,14 @@
 class UserContext:
     """Basic User Context Implementation"""
 
-    def __init__(self, buffsize: int) -> None:
-        self.buffer = bytearray([0xFF] * buffsize)
+    def __init__(self, buffsize: int = None, buffer: bytearray = None) -> None:
+        if buffer is not None:
+            self.buffer = buffer
+        elif buffsize is not None:
+            self.buffer = bytearray([0xFF] * buffsize)
+        else:
+            raise ValueError("Either buffsize or buffer must be provided")
+        self.in_size = len(self.buffer)
 
     def read(self, cfg: "LFSConfig", block: int, off: int, size: int) -> bytearray:
         """read data
@@ -91,6 +97,7 @@ def __init__(self, file_path: str, *, create: bool = False) -> None:
 
         self._path = file_path
         self._fh = open(file_path, mode)
+        self.in_size = os.path.getsize(file_path)
 
     def read(self, cfg: "LFSConfig", block: int, off: int, size: int) -> bytearray:
         logging.getLogger(__name__).debug("LFS Read : Block: %d, Offset: %d, Size=%d" % (block, off, size))
@@ -154,6 +161,7 @@ def __init__(self, disk_path: str) -> None:
         )
         if self.device == win32file.INVALID_HANDLE_VALUE:
             raise IOError("Could not open disk %s" % disk_path)
+        self.in_size = win32file.GetFileSize(self.device)
 
     def read(self, cfg: "LFSConfig", block: int, off: int, size: int) -> bytearray:
         """read data

src/littlefs/repl.py

@@ -19,7 +19,7 @@ def __init__(self, fs: LittleFS) -> None:
         """Initialize the shell with a LittleFS handle."""
         super().__init__()
         self._fs = fs
-        self._mounted = False
+        self._mounted = True
         self._cwd = "/"
 
     def onecmd(self, line: str) -> bool:

test/cli/test_create_and_extract.py

@@ -12,38 +12,41 @@ def test_create_and_extract(tmp_path):
     (source_dir / "file1.txt").write_text("hello world")
     (source_dir / "subdir").mkdir()
     (source_dir / "subdir" / "file2.txt").write_text("test content")
-    
+
     # Create filesystem image
     image_file = tmp_path / "test_image.bin"
     create_argv = [
         "littlefs",
         "create",
         str(source_dir),
         str(image_file),
-        "--block-size", "512",
-        "--fs-size", "64KB",
+        "--block-size",
+        "512",
+        "--fs-size",
+        "64KB",
     ]
     assert main(create_argv) == 0
     assert image_file.exists()
-    
+
     # Extract filesystem image
     extract_dir = tmp_path / "extracted"
     extract_argv = [
         "littlefs",
         "extract",
         str(image_file),
         str(extract_dir),
-        "--block-size", "512",
+        "--block-size",
+        "512",
     ]
     assert main(extract_argv) == 0
     assert extract_dir.exists()
-    
+
     # Compare directories
     comparison = filecmp.dircmp(source_dir, extract_dir)
     assert not comparison.diff_files
     assert not comparison.left_only
     assert not comparison.right_only
-    
+
     # Verify file contents
     assert (extract_dir / "file1.txt").read_text() == "hello world"
     assert (extract_dir / "subdir" / "file2.txt").read_text() == "test content"

test/cli/test_create_and_repl.py

@@ -0,0 +1,99 @@
+from pathlib import Path
+from unittest.mock import patch
+from io import StringIO
+
+from littlefs.__main__ import main
+
+
+def test_create_compact_no_pad_and_repl(tmp_path):
+    """Test creating a filesystem image with --compact --no-pad and opening it in REPL."""
+    # Create test directory with files
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    (source_dir / "file1.txt").write_text("hello world")
+    (source_dir / "subdir").mkdir()
+    (source_dir / "subdir" / "file2.txt").write_text("test content")
+
+    # Create filesystem image with --compact and --no-pad flags
+    image_file = tmp_path / "test_compact.bin"
+    create_argv = [
+        "littlefs",
+        "create",
+        str(source_dir),
+        str(image_file),
+        "--block-size",
+        "512",
+        "--fs-size",
+        "64KB",
+        "--compact",
+        "--no-pad",
+    ]
+    assert main(create_argv) == 0
+    assert image_file.exists()
+
+    # Verify the image is compacted (size should be less than 64KB)
+    image_size = image_file.stat().st_size
+    assert image_size < 64 * 1024, f"Expected compacted size < 64KB, got {image_size}"
+
+    # Mock stdin to exit immediately from REPL
+    # The REPL will run cmdloop() which reads from stdin
+    # We send "exit" command to quit the REPL
+    mock_stdin = StringIO("exit\n")
+
+    with patch("sys.stdin", mock_stdin):
+        # Test that REPL can open and mount the compacted image
+        repl_argv = [
+            "littlefs",
+            "repl",
+            str(image_file),
+            "--block-size",
+            "512",
+        ]
+        # The REPL should successfully mount and then exit
+        result = main(repl_argv)
+        assert result == 0 or result is None  # REPL returns 0 or None on success
+
+
+def test_create_verbose_and_repl_verbose(tmp_path):
+    """Test creating a filesystem image with --verbose and opening it in REPL with --verbose."""
+    # Create test directory with files
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    (source_dir / "file1.txt").write_text("hello world")
+    (source_dir / "subdir").mkdir()
+    (source_dir / "subdir" / "file2.txt").write_text("test content")
+
+    # Create filesystem image with --verbose flag
+    image_file = tmp_path / "test_verbose.bin"
+
+    create_argv = [
+        "littlefs",
+        "create",
+        str(source_dir),
+        str(image_file),
+        "--block-size",
+        "512",
+        "--fs-size",
+        "64KB",
+        "--verbose",
+    ]
+    assert main(create_argv) == 0
+
+    assert image_file.exists()
+
+    # Mock stdin to exit immediately from REPL
+    mock_stdin = StringIO("exit\n")
+
+    with patch("sys.stdin", mock_stdin):
+        # Test that REPL can open and mount the image with --verbose
+        repl_argv = [
+            "littlefs",
+            "repl",
+            str(image_file),
+            "--block-size",
+            "512",
+            "--verbose",
+        ]
+        # The REPL should successfully mount and then exit
+        result = main(repl_argv)
+        assert result == 0 or result is None  # REPL returns 0 or None on success