kythe.io@v0.0.68-0.20240422202219-7225dbc01741/kythe/go/util/tools/vnames/apply.go (about)

     1  /*
     2   * Copyright 2020 The Kythe Authors. All rights reserved.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *   http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package main
    18  
    19  import (
    20  	"bufio"
    21  	"context"
    22  	"encoding/json"
    23  	"flag"
    24  	"fmt"
    25  	"io"
    26  	"io/ioutil"
    27  	"os"
    28  	"strings"
    29  
    30  	"kythe.io/kythe/go/platform/vfs"
    31  	"kythe.io/kythe/go/util/cmdutil"
    32  	"kythe.io/kythe/go/util/log"
    33  	"kythe.io/kythe/go/util/vnameutil"
    34  
    35  	"github.com/google/subcommands"
    36  	"google.golang.org/protobuf/encoding/protojson"
    37  )
    38  
    39  type rulesFormat string
    40  
    41  const (
    42  	jsonFormat  rulesFormat = "JSON"
    43  	protoFormat rulesFormat = "PROTO"
    44  )
    45  
    46  type applyRulesCmd struct {
    47  	cmdutil.Info
    48  
    49  	rulesPath string
    50  	format    string
    51  
    52  	allowEmpty bool
    53  }
    54  
    55  var applyRulesInfo = cmdutil.NewInfo("apply-rules", "apply VName rewrite rules to stdin lines",
    56  	`Usage: apply-rules --rules <path> [--format proto]`)
    57  
    58  func (c *applyRulesCmd) SetFlags(flag *flag.FlagSet) {
    59  	flag.StringVar(&c.rulesPath, "rules", "", "Path to VName rewrite rules file")
    60  	flag.StringVar(&c.format, "format", string(jsonFormat), `Path of VName rewrite rules file {"JSON", "PROTO"}`)
    61  	flag.BoolVar(&c.allowEmpty, "allow_empty", false, "Whether to write an empty VName when no rule matches the input")
    62  }
    63  func (c *applyRulesCmd) Execute(ctx context.Context, flag *flag.FlagSet, args ...any) subcommands.ExitStatus {
    64  	rules, err := rulesFormat(strings.ToUpper(c.format)).readFile(ctx, c.rulesPath)
    65  	if err != nil {
    66  		return cmdErrorf("reading %q: %v", c.rulesPath, err)
    67  	}
    68  	s := bufio.NewScanner(os.Stdin)
    69  	m := protojson.MarshalOptions{UseProtoNames: true, Indent: "  "}
    70  	for s.Scan() {
    71  		line := s.Text()
    72  		v, ok := rules.Apply(line)
    73  		if ok || c.allowEmpty {
    74  			rec, err := m.Marshal(v)
    75  			if err != nil {
    76  				return cmdErrorf("encoding VName: %v", err)
    77  			} else if _, err := os.Stdout.Write(rec); err != nil {
    78  				return cmdErrorf("writing VName: %v", err)
    79  			} else if _, err := fmt.Fprintln(os.Stdout, ""); err != nil {
    80  				return cmdErrorf("writing newline: %v", err)
    81  			}
    82  		}
    83  	}
    84  	if err := s.Err(); err != nil {
    85  		return cmdErrorf("reading paths from stdin: %v", err)
    86  	}
    87  	return subcommands.ExitSuccess
    88  }
    89  
    90  func (f rulesFormat) readFile(ctx context.Context, path string) (vnameutil.Rules, error) {
    91  	file, err := vfs.Open(ctx, path)
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  	defer file.Close()
    96  	return f.readRules(file)
    97  }
    98  
    99  func (f rulesFormat) readRules(r io.Reader) (vnameutil.Rules, error) {
   100  	switch f {
   101  	case jsonFormat:
   102  		return vnameutil.ReadRules(r)
   103  	case protoFormat:
   104  		rec, err := ioutil.ReadAll(r)
   105  		if err != nil {
   106  			return nil, err
   107  		}
   108  		return vnameutil.ParseProtoRules(rec)
   109  	default:
   110  		return nil, fmt.Errorf("unknown rules format: %q", f)
   111  	}
   112  }
   113  
   114  func (f rulesFormat) writeRules(rules vnameutil.Rules, w io.Writer) error {
   115  	switch f {
   116  	case jsonFormat:
   117  		en := json.NewEncoder(w)
   118  		en.SetIndent("", "  ")
   119  		return en.Encode(rules)
   120  	case protoFormat:
   121  		rec, err := rules.Marshal()
   122  		if err != nil {
   123  			return err
   124  		}
   125  		_, err = w.Write(rec)
   126  		return err
   127  	default:
   128  		return fmt.Errorf("unknown rules format: %q", f)
   129  	}
   130  }
   131  
   132  func cmdErrorf(msg string, args ...any) subcommands.ExitStatus {
   133  	log.Errorf(msg, args...)
   134  	return subcommands.ExitFailure
   135  }