From df380bca96bc60032e1d747aa014a806e43c63d2 Mon Sep 17 00:00:00 2001
From: Ashwin Naren <arihant2math@gmail.com>
Date: Thu, 8 May 2025 21:26:28 -0700
Subject: [PATCH] lzma FORMAT_ALONE implementation (#5777)

* implement init_alone

* error if format is raw and there is a memlimit
---
 stdlib/src/lzma.rs | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/stdlib/src/lzma.rs b/stdlib/src/lzma.rs
index ed8be9a69..21ba8b64c 100644
--- a/stdlib/src/lzma.rs
+++ b/stdlib/src/lzma.rs
@@ -39,7 +39,7 @@ mod _lzma {
     use rustpython_vm::types::Constructor;
     use rustpython_vm::{PyObjectRef, PyPayload, PyResult, VirtualMachine};
     use std::fmt;
-    use xz2::stream::{Action, Check, Error, Filters, Status, Stream};
+    use xz2::stream::{Action, Check, Error, Filters, LzmaOptions, Status, Stream};
 
     #[cfg(windows)]
     type EnumVal = i32;
@@ -149,6 +149,11 @@ mod _lzma {
         type Args = LZMADecompressorConstructorArgs;
 
         fn py_new(cls: PyTypeRef, args: Self::Args, vm: &VirtualMachine) -> PyResult {
+            if args.format == FORMAT_RAW && args.memlimit.is_some() {
+                return Err(
+                    vm.new_value_error("Cannot specify memory limit with FORMAT_RAW".to_string())
+                );
+            }
             let memlimit = args.memlimit.unwrap_or(u64::MAX);
             let filters = args.filters.unwrap_or(0);
             let stream_result = match args.format {
@@ -322,6 +327,25 @@ mod _lzma {
                     .map_err(|_| new_lzma_error("Failed to initialize encoder", vm))?)
             }
         }
+
+        fn init_alone(
+            preset: u32,
+            filter_specs: Option<Vec<PyObjectRef>>,
+            vm: &VirtualMachine,
+        ) -> PyResult<Stream> {
+            if let Some(_filter_specs) = filter_specs {
+                Err(new_lzma_error(
+                    "TODO: RUSTPYTHON: LZMA: Alone filter specs",
+                    vm,
+                ))
+            } else {
+                let options = LzmaOptions::new_preset(preset)
+                    .map_err(|_| new_lzma_error("Failed to initialize encoder", vm))?;
+                let stream = Stream::new_lzma_encoder(&options)
+                    .map_err(|_| new_lzma_error("Failed to initialize encoder", vm))?;
+                Ok(stream)
+            }
+        }
     }
 
     #[derive(FromArgs)]
@@ -365,7 +389,8 @@ mod _lzma {
             }
             let stream = match args.format {
                 FORMAT_XZ => Self::init_xz(args.check, preset, args.filters, vm)?,
-                // TODO: ALONE AND RAW
+                FORMAT_ALONE => Self::init_alone(preset, args.filters, vm)?,
+                // TODO: RAW
                 _ => return Err(new_lzma_error("Invalid format", vm)),
             };
             Ok(Self {