-
Notifications
You must be signed in to change notification settings - Fork 0
/
string.go
161 lines (143 loc) · 5 KB
/
string.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
// Package string defines functions that manipulate strings, it's intended to be a drop-in subset of Python's string module for Starlark.
// See https://docs.python.org/3/library/string.html and https://github.com/python/cpython/blob/main/Lib/string.py for reference.
package string
import (
"fmt"
"html"
"strconv"
"sync"
"unicode/utf8"
"go.starlark.net/starlark"
"go.starlark.net/starlarkstruct"
)
// ModuleName defines the expected name for this Module when used
// in starlark's load() function, eg: load('string', 'length')
const ModuleName = "string"
var (
once sync.Once
strModule starlark.StringDict
)
const (
whitespace = " \t\n\r\v\f"
asciiLowerCase = `abcdefghijklmnopqrstuvwxyz`
asciiUpperCase = `ABCDEFGHIJKLMNOPQRSTUVWXYZ`
asciiLetters = asciiLowerCase + asciiUpperCase
decDigits = `0123456789`
hexDigits = decDigits + `abcdefABCDEF`
octDigits = `01234567`
punctuation = `!"#$%&'()*+,-./:;<=>?@[\]^_` + "{|}~`"
printable = decDigits + asciiLetters + punctuation + whitespace
)
// LoadModule loads the string module. It is concurrency-safe and idempotent.
func LoadModule() (starlark.StringDict, error) {
once.Do(func() {
strModule = starlark.StringDict{
ModuleName: &starlarkstruct.Module{
Name: ModuleName,
Members: starlark.StringDict{
// constants
"ascii_lowercase": starlark.String(asciiLowerCase),
"ascii_uppercase": starlark.String(asciiUpperCase),
"ascii_letters": starlark.String(asciiLetters),
"digits": starlark.String(decDigits),
"hexdigits": starlark.String(hexDigits),
"octdigits": starlark.String(octDigits),
"punctuation": starlark.String(punctuation),
"whitespace": starlark.String(whitespace),
"printable": starlark.String(printable),
// functions
"length": starlark.NewBuiltin(ModuleName+".length", length),
"reverse": starlark.NewBuiltin(ModuleName+".reverse", reverse),
"escape": genStarStrBuiltin("escape", html.EscapeString),
"unescape": genStarStrBuiltin("unescape", html.UnescapeString),
"quote": genStarStrBuiltin("quote", strconv.Quote),
"unquote": genStarStrBuiltin("unquote", robustUnquote),
},
},
}
})
return strModule, nil
}
// for convenience
var (
emptyStr string
none = starlark.None
)
type (
starFn func(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error)
)
// length returns the length of the given value.
func length(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
if l := len(args); l != 1 {
return none, fmt.Errorf(`length() takes exactly one argument (%d given)`, l)
}
switch r := args[0]; v := r.(type) {
case starlark.String:
return starlark.MakeInt(utf8.RuneCountInString(v.GoString())), nil
case starlark.Bytes:
return starlark.MakeInt(len(v)), nil
default:
if sv, ok := v.(starlark.Sequence); ok {
return starlark.MakeInt(sv.Len()), nil
}
return none, fmt.Errorf(`length() function isn't supported for '%s' type object`, v.Type())
}
}
// reverse returns the reversed string of the given value.
func reverse(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
if l := len(args); l != 1 {
return none, fmt.Errorf(`reverse() takes exactly one argument (%d given)`, l)
}
switch r := args[0]; v := r.(type) {
case starlark.String:
rs := []rune(v.GoString())
for i, j := 0, len(rs)-1; i < j; i, j = i+1, j-1 {
rs[i], rs[j] = rs[j], rs[i]
}
return starlark.String(rs), nil
case starlark.Bytes:
bs := []byte(v)
for i, j := 0, len(v)-1; i < j; i, j = i+1, j-1 {
bs[i], bs[j] = bs[j], bs[i]
}
return starlark.Bytes(bs), nil
default:
return none, fmt.Errorf(`reverse() function isn't supported for '%s' type object`, v.Type())
}
}
// genStarStrBuiltin generates the string operation builtin for Starlark.
func genStarStrBuiltin(fn string, opFn func(string) string) *starlark.Builtin {
sf := func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
if l := len(args); l != 1 {
return none, fmt.Errorf(`%s() takes exactly one argument (%d given)`, fn, l)
}
switch r := args[0]; v := r.(type) {
case starlark.String:
return starlark.String(opFn(v.GoString())), nil
case starlark.Bytes:
return starlark.Bytes(opFn(string(v))), nil
default:
return none, fmt.Errorf(`%s() function isn't supported for '%s' type object`, fn, v.Type())
}
}
return starlark.NewBuiltin(ModuleName+"."+fn, sf)
}
// robustUnquote unquotes a string, even if it's not quoted.
func robustUnquote(s string) string {
if len(s) < 2 {
return s
}
// if it's not quoted, quote it
old := s
if !(s[0] == '"' && s[len(s)-1] == '"') {
s = `"` + s + `"`
}
// try to unquote
ns, err := strconv.Unquote(s)
if err != nil {
// if failed, return original string
return old
}
// return unmodified string
return ns
}