docs

README.md

@@ -25,23 +25,14 @@ vcfexpr filter -e "return variant:info('AN') > 3000" \
 ```
 
 ---
-use `all` in user-defined `funcs.lua`
-check the sample fields to get variants where all samples have high DP
+check the sample fields to get variants where `all` samples have high DP.
+`all` is defined by `vcfexpr` (`any`, `filter` are also available).
+Users can load their own functions with `-l $lua_file`.
 ```
-vcfexpr filter -l funcs.lua \
-   -e 'return all(variant:format("DP"), function (dp)  return dp > 10 end)' \
+vcfexpr filter \
+   -e 'return all(function (dp)  return dp > 10 end, variant:format("DP"))' \
    -o all-high-dp.bcf $input_vcf
 ```
-```
-$ cat funcs.lua
-function all(t, f)
-    for _, v in pairs(t) do
-        if v ~= nil and not f(v) then
-            return false
-        end
-    end
-return true
-```
 ---
 
 get all of the FORMAT fields for a single sample into a lua table.
@@ -143,18 +134,28 @@ Arguments:
   <PATH>  Path to input VCF or BCF file
 
 Options:
-  -e, --expression <EXPRESSION>    boolean Lua expression(s) to filter the VCF or BCF file
-  -t, --template <TEMPLATE>        template expression in luau: https://luau-lang.org/syntax#string-interpolation. e.g. '{variant.chrom}:{variant.pos}'
-  -l, --lua <LUA>                  File(s) containing lua code to load. May contain functions that will be called by the expressions
-  -p, --lua-prelude <LUA_PRELUDE>  File containing lua code to run once before any variants are processed
-  -o, --output <OUTPUT>            Optional output file. Default is stdout
-  -h, --help                       Print help
+  -e, --expression <EXPRESSION>
+          boolean Lua expression(s) to filter the VCF or BCF file
+  -s, --set-expression <SET_EXPRESSION>
+          expression(s) to set existing INFO fields (new ones can be added in prelude)
+          e.g. --set-expression "AFmax=math.max(variant:info('AF'), variant:info('AFx'))"
+  -t, --template <TEMPLATE>
+          template expression in luau: https://luau-lang.org/syntax#string-interpolation. e.g. '{variant.chrom}:{variant.pos}'
+  -l, --lua <LUA>
+          File(s) containing lua code to load. May contain functions that will be called by the expressions
+  -p, --lua-prelude <LUA_PRELUDE>
+          File containing lua code to run once before any variants are processed
+  -o, --output <OUTPUT>
+          Optional output file. Default is stdout
+  -h, --help
+          Print help
 ```
 
 
 # TODO
 
-+ Currently --info-expressions can only be used when output is VCF. Update to support template output as well. So we need the header to translate.
++ Currently --set-expressions can only be used when output is VCF. Update to support template output as well. So we need the header to translate.
++ suuport --set-expressions for FORMAT fields (the infrastructure for this is there, just have to expose it)
 + add a functional lib such as [Moses](https://github.com/Yonaba/Moses) or [Lume](https://github.com/rxi/lume) which have `map`/`filter` and other functions.
   (The user can add these on their own with `--lua`).
 + write a class to simplify accessing CSQ fields.

src/main.rs

@@ -26,10 +26,10 @@ pub enum Commands {
         #[arg(short, long)]
         expression: Vec<String>,
 
-        /// set expressions to set existing INFO fields (new ones can be added in prelude)
-        /// e.g. --info-expressions "AFmax=math.max(variant:info('AF'), variant:info('AFx'))"
+        /// expression(s) to set existing INFO field(s) (new ones can be added in prelude)
+        /// e.g. --set-expression "AFmax=math.max(variant:info('AF'), variant:info('AFx'))"
         #[arg(short = 's', long)]
-        info_expressions: Vec<String>,
+        set_expression: Vec<String>,
 
         /// template expression in luau: https://luau-lang.org/syntax#string-interpolation. e.g. '{variant.chrom}:{variant.pos}'
         #[arg(short, long)]
@@ -52,7 +52,7 @@ pub enum Commands {
 fn filter_main(
     path: String,
     expressions: Vec<String>,
-    info_expressions: Vec<String>,
+    set_expression: Vec<String>,
     template: Option<String>,
     lua_code: Vec<String>,
     lua_prelude: Option<String>,
@@ -65,7 +65,7 @@ fn filter_main(
         &lua,
         path,
         expressions,
-        info_expressions,
+        set_expression,
         template,
         lua_prelude,
         output,
@@ -91,7 +91,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         Some(Commands::Filter {
             path,
             expression,
-            info_expressions,
+            set_expression,
             template,
             lua: lua_code,
             lua_prelude,
@@ -100,7 +100,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             filter_main(
                 path,
                 expression,
-                info_expressions,
+                set_expression,
                 template,
                 lua_code,
                 lua_prelude,

src/pprint.rs

@@ -86,3 +86,40 @@ function pprint(node)
     print(output_str)
 end
         "#;
+
+pub const PRELUDE: &str = r#"
+function map(f, t, skip_nil)
+    local new_t = {}
+    local j = 1
+    for i, v in ipairs(t) do
+        if v ~= nil or not skip_nil then
+            new_t[j] = f(v)
+            j = j + 1
+        end
+    end
+    return new_t
+end
+function filter(f, t, skip_nil)
+    local new_t = {}
+    local j = 1
+    for i, v in ipairs(t) do
+        if v ~= nil or not skip_nil then
+            if f(v) then
+                new_t[j] = v
+                j = j + 1
+            end
+        end
+    end
+    return new_t
+end
+
+function all(f, t, skip_nil)
+    for i, v in ipairs(t) do
+        if (v ~= nil or not skip_nil) and not f(v) then
+            return false
+        end
+    end
+    return true
+end
+
+"#;

src/vcfexpr.rs

@@ -14,7 +14,7 @@ pub struct VCFExpr<'lua> {
     template: Option<mlua::Function<'lua>>,
     writer: Option<EitherWriter>,
     expressions: Vec<mlua::Function<'lua>>,
-    info_expressions: HashMap<InfoFormat, ((TagType, TagLength), mlua::Function<'lua>)>,
+    set_expressions: HashMap<InfoFormat, ((TagType, TagLength), mlua::Function<'lua>)>,
     globals: mlua::Table<'lua>,
     variants_evaluated: usize,
     variants_passing: usize,
@@ -124,12 +124,15 @@ impl<'lua> VCFExpr<'lua> {
         lua: &'lua Lua,
         vcf_path: String,
         expression: Vec<String>,
-        info_expressions: Vec<String>,
+        set_expression: Vec<String>,
         template: Option<String>,
         lua_prelude: Option<String>,
         output: Option<String>,
     ) -> Result<Self, Box<dyn std::error::Error>> {
         lua.load(crate::pprint::PPRINT).set_name("pprint").exec()?;
+        lua.load(crate::pprint::PRELUDE)
+            .set_name("prelude")
+            .exec()?;
 
         let mut reader = match vcf_path.as_str() {
             "-" | "stdin" => bcf::Reader::from_stdin()?,
@@ -162,7 +165,7 @@ impl<'lua> VCFExpr<'lua> {
             })?;
         }
 
-        let info_exps = VCFExpr::load_info_expressions(lua, &mut hv, info_expressions)?;
+        let info_exps = VCFExpr::load_info_expressions(lua, &mut hv, set_expression)?;
 
         let header = bcf::header::Header::from_template(&hv);
 
@@ -189,7 +192,7 @@ impl<'lua> VCFExpr<'lua> {
             template,
             writer: Some(writer),
             expressions: exps,
-            info_expressions: info_exps,
+            set_expressions: info_exps,
             globals,
             variants_evaluated: 0,
             variants_passing: 0,
@@ -258,7 +261,7 @@ impl<'lua> VCFExpr<'lua> {
         &self,
         info_results: &mut HashMap<String, InfoFormatValue>,
     ) -> mlua::Result<()> {
-        for (inf, ((tagtyp, _taglen), expr)) in self.info_expressions.iter() {
+        for (inf, ((tagtyp, _taglen), expr)) in self.set_expressions.iter() {
             if let InfoFormat::Info(tag) = inf {
                 let t = match tagtyp {
                     TagType::Flag => {