go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/projects/nodes/pkg/funcs/filter_table_expression.go (about)

     1  /*
     2  
     3  Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.
     5  
     6  */
     7  
     8  package funcs
     9  
    10  import (
    11  	"context"
    12  	"fmt"
    13  
    14  	"github.com/wcharczuk/go-incr"
    15  	"go.charczuk.com/projects/nodes/pkg/incrutil"
    16  	"go.charczuk.com/projects/nodes/pkg/types"
    17  	"go.starlark.net/starlark"
    18  )
    19  
    20  // FilterTableExpression applies a predicate to a table on a row by row basis.
    21  func FilterTableExpression(expression string) func(context.Context, *types.Table, *incrutil.Inputs[any]) (*types.Table, error) {
    22  	return func(ctx context.Context, table *types.Table, inputs *incrutil.Inputs[any]) (output *types.Table, err error) {
    23  		if table == nil {
    24  			return
    25  		}
    26  		output = new(types.Table)
    27  
    28  		var row int
    29  		columnfn := func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
    30  			var n int = 0
    31  			if err := starlark.UnpackArgs(b.Name(), args, kwargs, "n", &n); err != nil {
    32  				return nil, err
    33  			}
    34  			if n < 0 || n > len(table.Columns) {
    35  				return nil, fmt.Errorf("input: cannot access invalid column index: %d", n)
    36  			}
    37  			return castAsStarlarkValue(table.Columns[n].Values[row])
    38  		}
    39  
    40  		var numRows int
    41  		for _, col := range table.Columns {
    42  			if len(col.Values) > numRows {
    43  				numRows = len(col.Values)
    44  			}
    45  		}
    46  
    47  		inputFunc := createInputFunc[any](inputs)
    48  		var builtins = starlark.StringDict{
    49  			"column": starlark.NewBuiltin("column", columnfn),
    50  			"input":  starlark.NewBuiltin("input", inputFunc),
    51  		}
    52  		for k, v := range defaultBuiltins {
    53  			builtins[k] = v
    54  		}
    55  
    56  		thread := new(starlark.Thread)
    57  		if tracer := incr.GetTracer(ctx); tracer != nil {
    58  			thread.Print = func(_ *starlark.Thread, msg string) {
    59  				tracer.Print("filter_table_expression:", msg)
    60  			}
    61  		}
    62  
    63  		output.Columns = make([]types.TableColumn, len(table.Columns))
    64  		for x := 0; x < len(output.Columns); x++ {
    65  			output.Columns[x].Name = table.Columns[x].Name
    66  			output.Columns[x].ValueType = table.Columns[x].ValueType
    67  		}
    68  
    69  		_, compiledFunction, err := starlark.SourceProgramOptions(&expressionFileOptions, "<expr>", expression, builtins.Has)
    70  		if err != nil {
    71  			return nil, err
    72  		}
    73  
    74  		for ; row < numRows; row++ {
    75  			var sd starlark.StringDict
    76  			sd, err = compiledFunction.Init(thread, builtins)
    77  			if err != nil {
    78  				return
    79  			}
    80  			if sd != nil {
    81  				sd.Freeze()
    82  				v, ok := sd["output"]
    83  				if !ok {
    84  					err = fmt.Errorf("filter table expression: you must use the `output` global to set the output")
    85  					return
    86  				}
    87  				switch typedv := v.(type) {
    88  				case starlark.Bool:
    89  					if typedv.Truth() == starlark.True {
    90  						for x := 0; x < len(output.Columns); x++ {
    91  							output.Columns[x].Values = append(output.Columns[x].Values, table.Columns[x].Values[row])
    92  						}
    93  					}
    94  				}
    95  			}
    96  		}
    97  		return
    98  	}
    99  }