github.com/k14s/starlark-go@v0.0.0-20200720175618-3a5c849cc368/cmd/starlark/starlark.go (about)

     1  // Copyright 2017 The Bazel Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // The starlark command interprets a Starlark file.
     6  // With no arguments, it starts a read-eval-print loop (REPL).
     7  package main // import "github.com/k14s/starlark-go/cmd/starlark"
     8  
     9  import (
    10  	"flag"
    11  	"fmt"
    12  	"log"
    13  	"os"
    14  	"runtime"
    15  	"runtime/pprof"
    16  	"strings"
    17  
    18  	"github.com/k14s/starlark-go/internal/compile"
    19  	"github.com/k14s/starlark-go/repl"
    20  	"github.com/k14s/starlark-go/resolve"
    21  	"github.com/k14s/starlark-go/starlark"
    22  )
    23  
    24  // flags
    25  var (
    26  	cpuprofile = flag.String("cpuprofile", "", "gather Go CPU profile in this file")
    27  	memprofile = flag.String("memprofile", "", "gather Go memory profile in this file")
    28  	profile    = flag.String("profile", "", "gather Starlark time profile in this file")
    29  	showenv    = flag.Bool("showenv", false, "on success, print final global environment")
    30  	execprog   = flag.String("c", "", "execute program `prog`")
    31  )
    32  
    33  func init() {
    34  	flag.BoolVar(&compile.Disassemble, "disassemble", compile.Disassemble, "show disassembly during compilation of each function")
    35  
    36  	// non-standard dialect flags
    37  	flag.BoolVar(&resolve.AllowFloat, "float", resolve.AllowFloat, "allow floating-point numbers")
    38  	flag.BoolVar(&resolve.AllowSet, "set", resolve.AllowSet, "allow set data type")
    39  	flag.BoolVar(&resolve.AllowLambda, "lambda", resolve.AllowLambda, "allow lambda expressions")
    40  	flag.BoolVar(&resolve.AllowNestedDef, "nesteddef", resolve.AllowNestedDef, "allow nested def statements")
    41  	flag.BoolVar(&resolve.AllowRecursion, "recursion", resolve.AllowRecursion, "allow while statements and recursive functions")
    42  	flag.BoolVar(&resolve.AllowGlobalReassign, "globalreassign", resolve.AllowGlobalReassign, "allow reassignment of globals, and if/for/while statements at top level")
    43  }
    44  
    45  func main() {
    46  	os.Exit(doMain())
    47  }
    48  
    49  func doMain() int {
    50  	log.SetPrefix("starlark: ")
    51  	log.SetFlags(0)
    52  	flag.Parse()
    53  
    54  	if *cpuprofile != "" {
    55  		f, err := os.Create(*cpuprofile)
    56  		check(err)
    57  		err = pprof.StartCPUProfile(f)
    58  		check(err)
    59  		defer func() {
    60  			pprof.StopCPUProfile()
    61  			err := f.Close()
    62  			check(err)
    63  		}()
    64  	}
    65  	if *memprofile != "" {
    66  		f, err := os.Create(*memprofile)
    67  		check(err)
    68  		defer func() {
    69  			runtime.GC()
    70  			err := pprof.Lookup("heap").WriteTo(f, 0)
    71  			check(err)
    72  			err = f.Close()
    73  			check(err)
    74  		}()
    75  	}
    76  
    77  	if *profile != "" {
    78  		f, err := os.Create(*profile)
    79  		check(err)
    80  		err = starlark.StartProfile(f)
    81  		check(err)
    82  		defer func() {
    83  			err := starlark.StopProfile()
    84  			check(err)
    85  		}()
    86  	}
    87  
    88  	thread := &starlark.Thread{Load: repl.MakeLoad()}
    89  	globals := make(starlark.StringDict)
    90  
    91  	switch {
    92  	case flag.NArg() == 1 || *execprog != "":
    93  		var (
    94  			filename string
    95  			src      interface{}
    96  			err      error
    97  		)
    98  		if *execprog != "" {
    99  			// Execute provided program.
   100  			filename = "cmdline"
   101  			src = *execprog
   102  		} else {
   103  			// Execute specified file.
   104  			filename = flag.Arg(0)
   105  		}
   106  		thread.Name = "exec " + filename
   107  		globals, err = starlark.ExecFile(thread, filename, src, nil)
   108  		if err != nil {
   109  			repl.PrintError(err)
   110  			return 1
   111  		}
   112  	case flag.NArg() == 0:
   113  		fmt.Println("Welcome to Starlark (go.starlark.net)")
   114  		thread.Name = "REPL"
   115  		repl.REPL(thread, globals)
   116  		return 0
   117  	default:
   118  		log.Print("want at most one Starlark file name")
   119  		return 1
   120  	}
   121  
   122  	// Print the global environment.
   123  	if *showenv {
   124  		for _, name := range globals.Keys() {
   125  			if !strings.HasPrefix(name, "_") {
   126  				fmt.Fprintf(os.Stderr, "%s = %s\n", name, globals[name])
   127  			}
   128  		}
   129  	}
   130  
   131  	return 0
   132  }
   133  
   134  func check(err error) {
   135  	if err != nil {
   136  		log.Fatal(err)
   137  	}
   138  }