snail-lang 0.6.1__tar.gz → 0.6.2__tar.gz

{snail_lang-0.6.1 → snail_lang-0.6.2}/Cargo.lock

@@ -485,11 +485,11 @@ checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
 
 [[package]]
 name = "snail-ast"
-version = "0.6.1"
+version = "0.6.2"
 
 [[package]]
 name = "snail-core"
-version = "0.6.1"
+version = "0.6.2"
 dependencies = [
  "pyo3",
  "snail-ast",
@@ -500,14 +500,14 @@ dependencies = [
 
 [[package]]
 name = "snail-error"
-version = "0.6.1"
+version = "0.6.2"
 dependencies = [
  "snail-ast",
 ]
 
 [[package]]
 name = "snail-lower"
-version = "0.6.1"
+version = "0.6.2"
 dependencies = [
  "pyo3",
  "snail-ast",
@@ -516,7 +516,7 @@ dependencies = [
 
 [[package]]
 name = "snail-parser"
-version = "0.6.1"
+version = "0.6.2"
 dependencies = [
  "pest",
  "pest_derive",
@@ -526,7 +526,7 @@ dependencies = [
 
 [[package]]
 name = "snail-proptest"
-version = "0.6.1"
+version = "0.6.2"
 dependencies = [
  "proptest",
  "pyo3",
@@ -540,7 +540,7 @@ dependencies = [
 
 [[package]]
 name = "snail-python"
-version = "0.6.1"
+version = "0.6.2"
 dependencies = [
  "pyo3",
  "snail-core",

{snail_lang-0.6.1 → snail_lang-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snail-lang
-Version: 0.6.1
+Version: 0.6.2
 Requires-Dist: jmespath>=1.0.1
 Requires-Dist: maturin>=1.5 ; extra == 'dev'
 Requires-Dist: pytest ; extra == 'dev'
@@ -57,8 +57,10 @@ semicolons are optional. You can separate statements with newlines.
 Process files line-by-line with familiar awk semantics:
 
 ```snail-awk("hello world\nfoo bar\n")
+BEGIN { print("start") }
 /hello/ { print("matched:", $0) }
 { print($1, "->", $2) }
+END { print("done") }
 ```
 
 **Built-in variables:**
@@ -73,7 +75,11 @@ Process files line-by-line with familiar awk semantics:
 | `$p` | Current file path |
 | `$m` | Last regex match object |
 
-Begin/end blocks use CLI flags (`-b`/`--begin`, `-e`/`--end`) for setup and teardown:
+Begin/end blocks can live in the source file (`BEGIN { ... }` / `END { ... }`) or be supplied
+via CLI flags (`-b`/`--begin`, `-e`/`--end`) for setup and teardown. CLI BEGIN blocks run
+before in-file BEGIN blocks; CLI END blocks run after in-file END blocks.
+`BEGIN` and `END` are reserved keywords in all modes.
+BEGIN/END blocks are regular Snail blocks, so awk/map-only `$` variables are not available inside them.
 ```bash
 echo -e "5\n4\n3\n2\n1" | snail --awk --begin 'total = 0' --end 'print("Sum:", total)' '/^[0-9]+/ { total = total + int($1) }'
 ```
@@ -83,8 +89,10 @@ echo -e "5\n4\n3\n2\n1" | snail --awk --begin 'total = 0' --end 'print("Sum:", t
 Process files one at a time instead of line-by-line:
 
 ```snail-map
+BEGIN { print("start") }
 print("File:", $src)
 print("Size:", len($text), "bytes")
+END { print("done") }
 ```
 
 **Built-in variables:**
@@ -95,7 +103,11 @@ print("Size:", len($text), "bytes")
 | `$fd` | Open file handle for the current file |
 | `$text` | Lazy text view of the current file contents |
 
-Begin/end blocks use CLI flags (`-b`/`--begin`, `-e`/`--end`) for setup and teardown:
+Begin/end blocks can live in the source file (`BEGIN { ... }` / `END { ... }`) or be supplied
+via CLI flags (`-b`/`--begin`, `-e`/`--end`) for setup and teardown. CLI BEGIN blocks run
+before in-file BEGIN blocks; CLI END blocks run after in-file END blocks.
+BEGIN/END blocks are regular Snail blocks, so awk/map-only `$` variables are not available inside them.
+`BEGIN` and `END` are reserved keywords in all modes.
 ```bash
 snail --map --begin "print('start')" --end "print('done')" "print($src)" *.txt
 ```

{snail_lang-0.6.1 → snail_lang-0.6.2}/README.md

@@ -45,8 +45,10 @@ semicolons are optional. You can separate statements with newlines.
 Process files line-by-line with familiar awk semantics:
 
 ```snail-awk("hello world\nfoo bar\n")
+BEGIN { print("start") }
 /hello/ { print("matched:", $0) }
 { print($1, "->", $2) }
+END { print("done") }
 ```
 
 **Built-in variables:**
@@ -61,7 +63,11 @@ Process files line-by-line with familiar awk semantics:
 | `$p` | Current file path |
 | `$m` | Last regex match object |
 
-Begin/end blocks use CLI flags (`-b`/`--begin`, `-e`/`--end`) for setup and teardown:
+Begin/end blocks can live in the source file (`BEGIN { ... }` / `END { ... }`) or be supplied
+via CLI flags (`-b`/`--begin`, `-e`/`--end`) for setup and teardown. CLI BEGIN blocks run
+before in-file BEGIN blocks; CLI END blocks run after in-file END blocks.
+`BEGIN` and `END` are reserved keywords in all modes.
+BEGIN/END blocks are regular Snail blocks, so awk/map-only `$` variables are not available inside them.
 ```bash
 echo -e "5\n4\n3\n2\n1" | snail --awk --begin 'total = 0' --end 'print("Sum:", total)' '/^[0-9]+/ { total = total + int($1) }'
 ```
@@ -71,8 +77,10 @@ echo -e "5\n4\n3\n2\n1" | snail --awk --begin 'total = 0' --end 'print("Sum:", t
 Process files one at a time instead of line-by-line:
 
 ```snail-map
+BEGIN { print("start") }
 print("File:", $src)
 print("Size:", len($text), "bytes")
+END { print("done") }
 ```
 
 **Built-in variables:**
@@ -83,7 +91,11 @@ print("Size:", len($text), "bytes")
 | `$fd` | Open file handle for the current file |
 | `$text` | Lazy text view of the current file contents |
 
-Begin/end blocks use CLI flags (`-b`/`--begin`, `-e`/`--end`) for setup and teardown:
+Begin/end blocks can live in the source file (`BEGIN { ... }` / `END { ... }`) or be supplied
+via CLI flags (`-b`/`--begin`, `-e`/`--end`) for setup and teardown. CLI BEGIN blocks run
+before in-file BEGIN blocks; CLI END blocks run after in-file END blocks.
+BEGIN/END blocks are regular Snail blocks, so awk/map-only `$` variables are not available inside them.
+`BEGIN` and `END` are reserved keywords in all modes.
 ```bash
 snail --map --begin "print('start')" --end "print('done')" "print($src)" *.txt
 ```

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-ast/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "snail-ast"
-version = "0.6.1"
+version = "0.6.2"
 edition = "2024"
 readme = "README.md"
 

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "snail-core"
-version = "0.6.1"
+version = "0.6.2"
 edition.workspace = true
 readme = "README.md"
 

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-core/src/lib.rs

@@ -33,15 +33,22 @@ pub fn compile_snail_source_with_auto_print(
             Ok(module)
         }
         CompileMode::Map => {
-            let program = parse_map_program(source)?;
-            let module = lower_map_program_with_auto_print(py, &program, auto_print_last)?;
+            let (program, begin_blocks, end_blocks) = parse_map_program_with_begin_end(source)?;
+            let module = lower_map_program_with_begin_end(
+                py,
+                &program,
+                &begin_blocks,
+                &end_blocks,
+                auto_print_last,
+            )?;
             Ok(module)
         }
     }
 }
 
 /// Compile an awk program with separate begin and end code blocks.
-/// Each begin/end source is parsed as a regular Snail program.
+/// In-file BEGIN/END blocks are merged with CLI blocks so CLI BEGIN runs first
+/// and CLI END runs last. CLI blocks are parsed as regular Snail programs.
 pub fn compile_awk_source_with_begin_end(
     py: Python<'_>,
     main_source: &str,
@@ -55,7 +62,8 @@ pub fn compile_awk_source_with_begin_end(
 }
 
 /// Compile a map program with separate begin and end code blocks.
-/// Each begin/end source is parsed as a map program.
+/// In-file BEGIN/END blocks are merged with CLI blocks so CLI BEGIN runs first
+/// and CLI END runs last. CLI blocks are parsed as regular Snail programs.
 pub fn compile_map_source_with_begin_end(
     py: Python<'_>,
     main_source: &str,
@@ -63,17 +71,21 @@ pub fn compile_map_source_with_begin_end(
     end_sources: &[&str],
     auto_print_last: bool,
 ) -> Result<PyObject, SnailError> {
-    let program = parse_map_program(main_source)?;
-    let mut begin_blocks = Vec::new();
+    let (program, mut begin_blocks, mut end_blocks) =
+        parse_map_program_with_begin_end(main_source)?;
+
+    let mut cli_begin_blocks = Vec::new();
     for source in begin_sources {
-        let begin_program = parse_map_program(source)?;
+        let begin_program = parse_program(source)?;
         if !begin_program.stmts.is_empty() {
-            begin_blocks.push(begin_program.stmts);
+            cli_begin_blocks.push(begin_program.stmts);
         }
     }
-    let mut end_blocks = Vec::new();
+    cli_begin_blocks.extend(begin_blocks);
+    begin_blocks = cli_begin_blocks;
+
     for source in end_sources {
-        let end_program = parse_map_program(source)?;
+        let end_program = parse_program(source)?;
         if !end_program.stmts.is_empty() {
             end_blocks.push(end_program.stmts);
         }

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-error/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "snail-error"
-version = "0.6.1"
+version = "0.6.2"
 edition = "2024"
 readme = "README.md"
 

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-lower/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "snail-lower"
-version = "0.6.1"
+version = "0.6.2"
 edition = "2024"
 readme = "README.md"
 

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-parser/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "snail-parser"
-version = "0.6.1"
+version = "0.6.2"
 edition = "2024"
 readme = "README.md"
 

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-parser/src/lib.rs

@@ -12,13 +12,15 @@ mod string;
 mod util;
 
 use awk::parse_awk_rule;
-use stmt::parse_stmt_list;
-use util::{error_with_span, full_span, parse_error_from_pest};
+use stmt::{parse_block, parse_stmt, parse_stmt_list};
+use util::{error_with_span, full_span, parse_error_from_pest, span_from_offset, span_from_pair};
 
 #[derive(Parser)]
 #[grammar = "snail.pest"]
 pub struct SnailParser;
 
+pub type MapProgramWithBeginEnd = (Program, Vec<Vec<Stmt>>, Vec<Vec<Stmt>>);
+
 pub fn parse_program(source: &str) -> Result<Program, ParseError> {
     let mut pairs = SnailParser::parse(Rule::program, source)
         .map_err(|err| parse_error_from_pest(err, source))?;
@@ -45,29 +47,44 @@ pub fn parse_awk_program(source: &str) -> Result<AwkProgram, ParseError> {
         .ok_or_else(|| ParseError::new("missing awk program root"))?;
     let span = full_span(source);
 
+    let mut begin_blocks = Vec::new();
     let mut rules = Vec::new();
+    let mut end_blocks = Vec::new();
 
     for inner in pair.into_inner() {
         if inner.as_rule() == Rule::awk_entry_list {
             for entry in inner.into_inner() {
-                if entry.as_rule() == Rule::awk_rule {
-                    rules.push(parse_awk_rule(entry, source)?);
+                match entry.as_rule() {
+                    Rule::awk_begin => {
+                        let block = parse_begin_end_block(entry, source, "BEGIN")?;
+                        if !block.is_empty() {
+                            begin_blocks.push(block);
+                        }
+                    }
+                    Rule::awk_end => {
+                        let block = parse_begin_end_block(entry, source, "END")?;
+                        if !block.is_empty() {
+                            end_blocks.push(block);
+                        }
+                    }
+                    Rule::awk_rule => rules.push(parse_awk_rule(entry, source)?),
+                    _ => {}
                 }
             }
         }
     }
 
     Ok(AwkProgram {
-        begin_blocks: Vec::new(),
+        begin_blocks,
         rules,
-        end_blocks: Vec::new(),
+        end_blocks,
         span,
     })
 }
 
 /// Parses an awk program with separate begin and end code sources.
-/// Each begin/end source is parsed as a regular Snail program and its statements
-/// are injected into the resulting AwkProgram.
+/// Each begin/end source is parsed as a regular Snail program and merged so CLI BEGIN
+/// blocks run before in-file BEGIN blocks, and CLI END blocks run after in-file END blocks.
 pub fn parse_awk_program_with_begin_end(
     main_source: &str,
     begin_sources: &[&str],
@@ -75,21 +92,24 @@ pub fn parse_awk_program_with_begin_end(
 ) -> Result<AwkProgram, ParseError> {
     let mut program = parse_awk_program(main_source)?;
 
-    // Parse each begin source as a regular program and extract statements
+    let mut cli_begin_blocks = Vec::new();
     for source in begin_sources {
         let begin_program = parse_program(source)?;
         if !begin_program.stmts.is_empty() {
-            program.begin_blocks.push(begin_program.stmts);
+            cli_begin_blocks.push(begin_program.stmts);
         }
     }
+    cli_begin_blocks.extend(program.begin_blocks);
+    program.begin_blocks = cli_begin_blocks;
 
-    // Parse each end source as a regular program and extract statements
+    let mut end_blocks = program.end_blocks;
     for source in end_sources {
         let end_program = parse_program(source)?;
         if !end_program.stmts.is_empty() {
-            program.end_blocks.push(end_program.stmts);
+            end_blocks.push(end_program.stmts);
         }
     }
+    program.end_blocks = end_blocks;
 
     Ok(program)
 }
@@ -102,27 +122,133 @@ const MAP_ONLY_MESSAGE: &str = "map variables are only valid in map mode; use --
 
 /// Parses a map program that processes files one at a time.
 /// Allows map variables ($src, $fd, $text) but rejects awk variables.
+/// In-file BEGIN/END blocks are validated but not returned; use
+/// `parse_map_program_with_begin_end` to access them.
 pub fn parse_map_program(source: &str) -> Result<Program, ParseError> {
-    let mut pairs = SnailParser::parse(Rule::program, source)
+    let (program, _, _) = parse_map_program_with_begin_end(source)?;
+    Ok(program)
+}
+
+fn validate_no_awk_syntax_for_map(program: &Program, source: &str) -> Result<(), ParseError> {
+    for stmt in &program.stmts {
+        validate_stmt_for_map(stmt, source)?;
+    }
+    Ok(())
+}
+
+/// Parses a map program with in-file BEGIN/END blocks.
+/// BEGIN/END blocks are parsed as regular Snail statement blocks (no map/awk vars).
+pub fn parse_map_program_with_begin_end(
+    source: &str,
+) -> Result<MapProgramWithBeginEnd, ParseError> {
+    let mut pairs = SnailParser::parse(Rule::map_program, source)
         .map_err(|err| parse_error_from_pest(err, source))?;
     let pair = pairs
         .next()
-        .ok_or_else(|| ParseError::new("missing program root"))?;
+        .ok_or_else(|| ParseError::new("missing map program root"))?;
     let span = full_span(source);
     let mut stmts = Vec::new();
+    let mut begin_blocks = Vec::new();
+    let mut end_blocks = Vec::new();
+    let mut entries = Vec::new();
+
     for inner in pair.into_inner() {
-        if inner.as_rule() == Rule::stmt_list {
-            stmts = parse_stmt_list(inner, source)?;
+        if inner.as_rule() == Rule::map_entry_list {
+            for entry in inner.into_inner() {
+                if entry.as_rule() != Rule::map_entry {
+                    continue;
+                }
+                let entry_span = span_from_pair(&entry, source);
+                let mut entry_inner = entry.into_inner();
+                let entry_pair = entry_inner.next().ok_or_else(|| {
+                    error_with_span("missing map entry", entry_span.clone(), source)
+                })?;
+                match entry_pair.as_rule() {
+                    Rule::map_begin => {
+                        let block = parse_begin_end_block(entry_pair, source, "BEGIN")?;
+                        if !block.is_empty() {
+                            begin_blocks.push(block);
+                        }
+                        entries.push((entry_span, MapEntryKind::BeginEnd));
+                    }
+                    Rule::map_end => {
+                        let block = parse_begin_end_block(entry_pair, source, "END")?;
+                        if !block.is_empty() {
+                            end_blocks.push(block);
+                        }
+                        entries.push((entry_span, MapEntryKind::BeginEnd));
+                    }
+                    _ => {
+                        let stmt = parse_stmt(entry_pair, source)?;
+                        entries.push((entry_span, map_entry_kind_for_stmt(&stmt)));
+                        stmts.push(stmt);
+                    }
+                }
+            }
         }
     }
+
+    validate_map_entry_separators(&entries, source)?;
+
     let program = Program { stmts, span };
     validate_no_awk_syntax_for_map(&program, source)?;
-    Ok(program)
+    Ok((program, begin_blocks, end_blocks))
 }
 
-fn validate_no_awk_syntax_for_map(program: &Program, source: &str) -> Result<(), ParseError> {
-    for stmt in &program.stmts {
-        validate_stmt_for_map(stmt, source)?;
+fn parse_begin_end_block(
+    pair: pest::iterators::Pair<'_, Rule>,
+    source: &str,
+    label: &str,
+) -> Result<Vec<Stmt>, ParseError> {
+    let span = span_from_pair(&pair, source);
+    let mut inner = pair.into_inner();
+    let block_pair = inner
+        .next()
+        .ok_or_else(|| error_with_span(format!("missing {label} block"), span.clone(), source))?;
+    let block = parse_block(block_pair, source)?;
+    validate_block(&block, source)?;
+    Ok(block)
+}
+
+#[derive(Clone, Copy)]
+enum MapEntryKind {
+    BeginEnd,
+    Simple,
+    Compound,
+}
+
+fn map_entry_kind_for_stmt(stmt: &Stmt) -> MapEntryKind {
+    match stmt {
+        Stmt::If { .. }
+        | Stmt::While { .. }
+        | Stmt::For { .. }
+        | Stmt::Def { .. }
+        | Stmt::Class { .. }
+        | Stmt::Try { .. }
+        | Stmt::With { .. } => MapEntryKind::Compound,
+        _ => MapEntryKind::Simple,
+    }
+}
+
+fn validate_map_entry_separators(
+    entries: &[(SourceSpan, MapEntryKind)],
+    source: &str,
+) -> Result<(), ParseError> {
+    for window in entries.windows(2) {
+        let (prev_span, prev_kind) = &window[0];
+        let (next_span, next_kind) = &window[1];
+        let gap = &source[prev_span.end.offset..next_span.start.offset];
+        let has_sep = gap.contains('\n') || gap.contains(';');
+        if !has_sep
+            && matches!(prev_kind, MapEntryKind::Simple)
+            && !matches!(next_kind, MapEntryKind::BeginEnd)
+        {
+            return Err(error_with_span(
+                "expected statement separator",
+                span_from_offset(next_span.start.offset, next_span.start.offset, source),
+                source,
+            ));
+        }
     }
     Ok(())
 }

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-parser/src/snail.pest

@@ -1,13 +1,23 @@
 // Top-level program entry points
 program = { SOI ~ stmt_sep* ~ stmt_list? ~ stmt_sep* ~ EOI }
 awk_program = { SOI ~ stmt_sep* ~ awk_entry_list? ~ stmt_sep* ~ EOI }
+map_program = { SOI ~ stmt_sep* ~ map_entry_list? ~ stmt_sep* ~ EOI }
 
 // AWK mode: pattern-action rules
 awk_entry_list = { awk_entry ~ (stmt_sep* ~ awk_entry)* ~ stmt_sep* }
-awk_entry = _{ awk_rule }
+awk_entry = _{ awk_begin | awk_end | awk_rule }
+awk_begin = { "BEGIN" ~ block }
+awk_end = { "END" ~ block }
 awk_rule = { block | awk_pattern ~ block? }
 awk_pattern = { expr }
 
+// Map mode: program with optional BEGIN/END blocks
+map_entry_list = { map_entry ~ (stmt_sep* ~ map_entry)* ~ stmt_sep* }
+map_entry = { map_begin_end | stmt }
+map_begin_end = _{ map_begin | map_end }
+map_begin = { "BEGIN" ~ block }
+map_end = { "END" ~ block }
+
 // Statements: separated by semicolons or newlines
 // Compound statements (ending with }) don't need trailing separators
 // Simple statements need separators unless they're the last statement
@@ -261,6 +271,7 @@ keyword = _{
   | "return" | "break" | "continue" | "pass" | "raise" | "try" | "except"
   | "finally" | "with" | "assert" | "del" | "and" | "or" | "not" | "import" | "from" | "as"
   | "let"
+  | "BEGIN" | "END"
   | "True" | "False" | "None") ~ !ident_continue
 }
 

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-parser/tests/errors.rs

@@ -1,7 +1,10 @@
 mod common;
 
 use common::*;
-use snail_parser::{parse_awk_program, parse_awk_program_with_begin_end, parse_program};
+use snail_parser::{
+    parse_awk_program, parse_awk_program_with_begin_end, parse_map_program_with_begin_end,
+    parse_program,
+};
 
 #[test]
 fn reports_parse_error_with_location() {
@@ -203,20 +206,12 @@ fn parser_rejects_invalid_parameter_syntax() {
 // ========== AWK Mode Parser Tests ==========
 
 #[test]
-fn awk_begin_end_parsed_as_patterns() {
-    // BEGIN and END are no longer special keywords; they parse as regular pattern names
-    // Use -b/-e CLI flags for begin/end blocks instead
-    let program = parse_awk_program("BEGIN { print(1) }").expect("should parse");
+fn awk_begin_end_parsed_as_blocks() {
+    let program = parse_awk_program("BEGIN { print(1) } /foo/ { print($0) } END { print(2) }")
+        .expect("should parse");
     assert_eq!(program.rules.len(), 1);
-    assert!(program.begin_blocks.is_empty());
-    // BEGIN is parsed as pattern (identifier), not a special begin block
-    assert!(program.rules[0].pattern.is_some());
-
-    let program = parse_awk_program("END { print(1) }").expect("should parse");
-    assert_eq!(program.rules.len(), 1);
-    assert!(program.end_blocks.is_empty());
-    // END is parsed as pattern (identifier), not a special end block
-    assert!(program.rules[0].pattern.is_some());
+    assert_eq!(program.begin_blocks.len(), 1);
+    assert_eq!(program.end_blocks.len(), 1);
 }
 
 #[test]
@@ -246,6 +241,36 @@ fn awk_with_empty_begin_end() {
     assert!(program.end_blocks.is_empty());
 }
 
+#[test]
+fn awk_begin_end_rejects_awk_vars() {
+    let err = parse_awk_program("BEGIN { print($0) }").expect_err("should reject awk vars");
+    assert!(err.to_string().contains("$0"));
+}
+
+#[test]
+fn map_begin_end_parsed_as_blocks() {
+    let (program, begin_blocks, end_blocks) =
+        parse_map_program_with_begin_end("BEGIN { print(1) } print($src) END { print(2) }")
+            .expect("should parse");
+    assert_eq!(program.stmts.len(), 1);
+    assert_eq!(begin_blocks.len(), 1);
+    assert_eq!(end_blocks.len(), 1);
+}
+
+#[test]
+fn map_begin_end_rejects_map_vars() {
+    let err = parse_map_program_with_begin_end("BEGIN { print($src) }\nprint($src)")
+        .expect_err("should reject map vars in BEGIN/END");
+    assert!(err.to_string().contains("$src"));
+}
+
+#[test]
+fn map_requires_separators_between_simple_statements() {
+    let err = parse_map_program_with_begin_end("print($src) print($src)")
+        .expect_err("should reject missing separators");
+    assert!(err.to_string().contains("expected statement separator"));
+}
+
 // ========== F-String Interpolation Tests ==========
 
 #[test]

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "snail-python"
-version = "0.6.1"
+version = "0.6.2"
 edition.workspace = true
 build = "build.rs"
 

{snail_lang-0.6.1 → snail_lang-0.6.2}/crates/snail-python/src/lib.rs

@@ -5,9 +5,10 @@ use pyo3::exceptions::{PyRuntimeError, PySyntaxError, PySystemExit};
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyModule, PyTuple};
 use snail_core::{
-    CompileMode, ParseError, Program, compile_awk_source_with_begin_end,
+    CompileMode, ParseError, Program, Stmt, compile_awk_source_with_begin_end,
     compile_map_source_with_begin_end, compile_snail_source_with_auto_print, format_snail_error,
-    parse_awk_program, parse_awk_program_with_begin_end, parse_map_program, parse_program,
+    parse_awk_program, parse_awk_program_with_begin_end, parse_map_program_with_begin_end,
+    parse_program,
 };
 use std::sync::OnceLock;
 use std::time::Instant;
@@ -320,8 +321,8 @@ fn exec_py(
 #[derive(Debug)]
 struct MapAst {
     program: Program,
-    begin_blocks: Vec<Program>,
-    end_blocks: Vec<Program>,
+    begin_blocks: Vec<Vec<Stmt>>,
+    end_blocks: Vec<Vec<Stmt>>,
 }
 
 #[pyfunction(name = "parse_ast")]
@@ -351,24 +352,30 @@ fn parse_ast_py(
             Ok(format!("{:#?}", program))
         }
         CompileMode::Map => {
-            let program = parse_map_program(source).map_err(err_to_syntax)?;
-            if begin_code.is_empty() && end_code.is_empty() {
-                return Ok(format!("{:#?}", program));
-            }
-            let mut begin_blocks = Vec::new();
+            let (program, mut begin_blocks, mut end_blocks) =
+                parse_map_program_with_begin_end(source).map_err(err_to_syntax)?;
+
+            let mut cli_begin_blocks = Vec::new();
             for source in &begin_code {
-                let begin_program = parse_map_program(source).map_err(err_to_syntax)?;
+                let begin_program = parse_program(source).map_err(err_to_syntax)?;
                 if !begin_program.stmts.is_empty() {
-                    begin_blocks.push(begin_program);
+                    cli_begin_blocks.push(begin_program.stmts);
                 }
             }
-            let mut end_blocks = Vec::new();
+            cli_begin_blocks.extend(begin_blocks);
+            begin_blocks = cli_begin_blocks;
+
             for source in &end_code {
-                let end_program = parse_map_program(source).map_err(err_to_syntax)?;
+                let end_program = parse_program(source).map_err(err_to_syntax)?;
                 if !end_program.stmts.is_empty() {
-                    end_blocks.push(end_program);
+                    end_blocks.push(end_program.stmts);
                 }
             }
+
+            if begin_blocks.is_empty() && end_blocks.is_empty() {
+                return Ok(format!("{:#?}", program));
+            }
+
             let map_ast = MapAst {
                 program,
                 begin_blocks,
@@ -389,7 +396,7 @@ fn parse_py(source: &str, mode: &str, filename: &str) -> PyResult<()> {
         CompileMode::Awk => parse_awk_program(source)
             .map(|_| ())
             .map_err(|err| PySyntaxError::new_err(format_snail_error(&err.into(), filename))),
-        CompileMode::Map => parse_map_program(source)
+        CompileMode::Map => parse_map_program_with_begin_end(source)
             .map(|_| ())
             .map_err(|err| PySyntaxError::new_err(format_snail_error(&err.into(), filename))),
     }

{snail_lang-0.6.1 → snail_lang-0.6.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "snail-lang"
-version = "0.6.1"
+version = "0.6.2"
 description = "Snail programming language interpreter"
 readme = "README.md"
 requires-python = ">=3.10"

{snail_lang-0.6.1 → snail_lang-0.6.2}/python/snail/cli.py

@@ -187,13 +187,12 @@ def _parse_args(argv: list[str]) -> _Args:
             args.args = argv[idx + 1 :]
             return args
         if token == "-" or not token.startswith("-"):
-            if code_found:
-                # Already found code, rest are args
-                args.args.extend(argv[idx:])
-                return args
-            # This is the code, continue parsing for -b/--begin and -e/--end after
-            args.args = [token]
-            code_found = True
+            if not code_found:
+                # This is the code (or the first arg when -f is used)
+                args.args = [token]
+                code_found = True
+            else:
+                args.args.append(token)
             idx += 1
             continue
         if token in ("-h", "--help"):
@@ -231,6 +230,7 @@ def _parse_args(argv: list[str]) -> _Args:
             if idx + 1 >= len(argv):
                 raise ValueError("option -f requires an argument")
             args.file = argv[idx + 1]
+            code_found = True
             idx += 2
             continue
         if token in ("-b", "--begin"):