Fix bug in REPL and aggregate duplicated code

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 = len(fs.context.buffer)
+        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,41 +210,20 @@ 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})."
-        )
+    fs = _mount_from_context(parser, args, context)
 
-    block_count = image_size // args.block_size
-    if block_count == 0:
-        parser.error("Image is smaller than a single block; cannot mount.")
+    shell = LittleFSRepl(fs)
 
-    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.cmdloop()
 
-    shell = LittleFSRepl(fs)
-    try:
-        try:
-            shell.do_mount()
-        except LittleFSError as exc:
-            parser.error(f"Failed to mount '{source}': {exc}")
-        shell.cmdloop()
-    finally:
-        if shell._mounted:
-            with suppress(LittleFSError):
-                fs.unmount()
+    if shell._mounted:
+        with suppress(LittleFSError):
+            fs.unmount()
 
     return 0
 

src/littlefs/context.py

@@ -10,8 +10,13 @@
 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")
 
     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_repl.py

@@ -0,0 +1,51 @@
+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