github.com/neatlab/neatio@v1.7.3-0.20220425043230-d903e92fcc75/utilities/utils/customflags.go (about)

     1  package utils
     2  
     3  import (
     4  	"encoding"
     5  	"errors"
     6  	"flag"
     7  	"fmt"
     8  	"math/big"
     9  	"os"
    10  	"os/user"
    11  	"path"
    12  	"strings"
    13  
    14  	"github.com/neatlab/neatio/utilities/common/math"
    15  	"gopkg.in/urfave/cli.v1"
    16  )
    17  
    18  type DirectoryString struct {
    19  	Value string
    20  }
    21  
    22  func (self *DirectoryString) String() string {
    23  	return self.Value
    24  }
    25  
    26  func (self *DirectoryString) Set(value string) error {
    27  	self.Value = expandPath(value)
    28  	return nil
    29  }
    30  
    31  type DirectoryFlag struct {
    32  	Name  string
    33  	Value DirectoryString
    34  	Usage string
    35  }
    36  
    37  func (self DirectoryFlag) String() string {
    38  	fmtString := "%s %v\t%v"
    39  	if len(self.Value.Value) > 0 {
    40  		fmtString = "%s \"%v\"\t%v"
    41  	}
    42  	return fmt.Sprintf(fmtString, prefixedNames(self.Name), self.Value.Value, self.Usage)
    43  }
    44  
    45  func eachName(longName string, fn func(string)) {
    46  	parts := strings.Split(longName, ",")
    47  	for _, name := range parts {
    48  		name = strings.Trim(name, " ")
    49  		fn(name)
    50  	}
    51  }
    52  
    53  func (self DirectoryFlag) Apply(set *flag.FlagSet) {
    54  	eachName(self.Name, func(name string) {
    55  		set.Var(&self.Value, self.Name, self.Usage)
    56  	})
    57  }
    58  
    59  type TextMarshaler interface {
    60  	encoding.TextMarshaler
    61  	encoding.TextUnmarshaler
    62  }
    63  
    64  type textMarshalerVal struct {
    65  	v TextMarshaler
    66  }
    67  
    68  func (v textMarshalerVal) String() string {
    69  	if v.v == nil {
    70  		return ""
    71  	}
    72  	text, _ := v.v.MarshalText()
    73  	return string(text)
    74  }
    75  
    76  func (v textMarshalerVal) Set(s string) error {
    77  	return v.v.UnmarshalText([]byte(s))
    78  }
    79  
    80  type TextMarshalerFlag struct {
    81  	Name  string
    82  	Value TextMarshaler
    83  	Usage string
    84  }
    85  
    86  func (f TextMarshalerFlag) GetName() string {
    87  	return f.Name
    88  }
    89  
    90  func (f TextMarshalerFlag) String() string {
    91  	return fmt.Sprintf("%s \"%v\"\t%v", prefixedNames(f.Name), f.Value, f.Usage)
    92  }
    93  
    94  func (f TextMarshalerFlag) Apply(set *flag.FlagSet) {
    95  	eachName(f.Name, func(name string) {
    96  		set.Var(textMarshalerVal{f.Value}, f.Name, f.Usage)
    97  	})
    98  }
    99  
   100  func GlobalTextMarshaler(ctx *cli.Context, name string) TextMarshaler {
   101  	val := ctx.GlobalGeneric(name)
   102  	if val == nil {
   103  		return nil
   104  	}
   105  	return val.(textMarshalerVal).v
   106  }
   107  
   108  type BigFlag struct {
   109  	Name  string
   110  	Value *big.Int
   111  	Usage string
   112  }
   113  
   114  type bigValue big.Int
   115  
   116  func (b *bigValue) String() string {
   117  	if b == nil {
   118  		return ""
   119  	}
   120  	return (*big.Int)(b).String()
   121  }
   122  
   123  func (b *bigValue) Set(s string) error {
   124  	int, ok := math.ParseBig256(s)
   125  	if !ok {
   126  		return errors.New("invalid integer syntax")
   127  	}
   128  	*b = (bigValue)(*int)
   129  	return nil
   130  }
   131  
   132  func (f BigFlag) GetName() string {
   133  	return f.Name
   134  }
   135  
   136  func (f BigFlag) String() string {
   137  	fmtString := "%s %v\t%v"
   138  	if f.Value != nil {
   139  		fmtString = "%s \"%v\"\t%v"
   140  	}
   141  	return fmt.Sprintf(fmtString, prefixedNames(f.Name), f.Value, f.Usage)
   142  }
   143  
   144  func (f BigFlag) Apply(set *flag.FlagSet) {
   145  	eachName(f.Name, func(name string) {
   146  		set.Var((*bigValue)(f.Value), f.Name, f.Usage)
   147  	})
   148  }
   149  
   150  func GlobalBig(ctx *cli.Context, name string) *big.Int {
   151  	val := ctx.GlobalGeneric(name)
   152  	if val == nil {
   153  		return nil
   154  	}
   155  	return (*big.Int)(val.(*bigValue))
   156  }
   157  
   158  func prefixFor(name string) (prefix string) {
   159  	if len(name) == 1 {
   160  		prefix = "-"
   161  	} else {
   162  		prefix = "--"
   163  	}
   164  
   165  	return
   166  }
   167  
   168  func prefixedNames(fullName string) (prefixed string) {
   169  	parts := strings.Split(fullName, ",")
   170  	for i, name := range parts {
   171  		name = strings.Trim(name, " ")
   172  		prefixed += prefixFor(name) + name
   173  		if i < len(parts)-1 {
   174  			prefixed += ", "
   175  		}
   176  	}
   177  	return
   178  }
   179  
   180  func (self DirectoryFlag) GetName() string {
   181  	return self.Name
   182  }
   183  
   184  func (self *DirectoryFlag) Set(value string) {
   185  	self.Value.Value = value
   186  }
   187  
   188  func expandPath(p string) string {
   189  	if strings.HasPrefix(p, "~/") || strings.HasPrefix(p, "~\\") {
   190  		if home := homeDir(); home != "" {
   191  			p = home + p[1:]
   192  		}
   193  	}
   194  	return path.Clean(os.ExpandEnv(p))
   195  }
   196  
   197  func homeDir() string {
   198  	if home := os.Getenv("HOME"); home != "" {
   199  		return home
   200  	}
   201  	if usr, err := user.Current(); err == nil {
   202  		return usr.HomeDir
   203  	}
   204  	return ""
   205  }