github.com/galaxyobe/gen@v0.0.0-20220910125335-392fa8f0990f/pkg/util/method.go (about)

     1  /*
     2   Copyright 2022 Galaxyobe.
     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 util
    18  
    19  import (
    20  	"strconv"
    21  	"strings"
    22  
    23  	"github.com/iancoleman/strcase"
    24  	"k8s.io/gengo/types"
    25  )
    26  
    27  type MethodSet map[string][]types.Member // key: method name, value: types.Member
    28  
    29  func NewMethodSet() MethodSet {
    30  	return make(MethodSet)
    31  }
    32  
    33  func (m MethodSet) AddMethod(prefix, method string, member types.Member) {
    34  	if !strings.HasPrefix(method, prefix) {
    35  		method = prefix + method
    36  	}
    37  	list, ok := m[method]
    38  	if !ok {
    39  		m[method] = []types.Member{member}
    40  		return
    41  	}
    42  	list = append(list, member)
    43  	m[method] = list
    44  }
    45  
    46  func (m MethodSet) AddMethods(prefix string, methods []string, member types.Member) {
    47  	for _, method := range methods {
    48  		m.AddMethod(prefix, method, member)
    49  	}
    50  }
    51  
    52  type MethodGenerate struct {
    53  	set      map[string]int
    54  	gen      func(name string) string
    55  	handlers []func(string) string
    56  }
    57  
    58  func NewMethodGenerate(fn func(name string) string, handler ...func(string) string) MethodGenerate {
    59  	if len(handler) == 0 {
    60  		handler = append(handler, strcase.ToCamel)
    61  	}
    62  	return MethodGenerate{
    63  		set:      make(map[string]int),
    64  		gen:      fn,
    65  		handlers: handler,
    66  	}
    67  }
    68  
    69  func (m *MethodGenerate) handle(name string) string {
    70  	for _, f := range m.handlers {
    71  		name = f(name)
    72  	}
    73  	return name
    74  }
    75  
    76  func (m *MethodGenerate) AddExistName(name string) (bool, int) {
    77  	count, ok := m.set[name]
    78  	if ok {
    79  		count += 1
    80  	} else {
    81  		count = 1
    82  	}
    83  	m.set[name] = count
    84  	return ok, count
    85  }
    86  
    87  func (m *MethodGenerate) AddExistNames(name ...string) *MethodGenerate {
    88  	for _, item := range name {
    89  		m.AddExistName(item)
    90  	}
    91  	return m
    92  }
    93  
    94  func (m *MethodGenerate) ExistName(name string) bool {
    95  	_, ok := m.set[name]
    96  	return ok
    97  }
    98  
    99  func (m *MethodGenerate) GenName(name string) string {
   100  	name = m.handle(name)
   101  	method := m.gen(name)
   102  	if ok, count := m.AddExistName(method); ok {
   103  		method = m.gen(name + strconv.Itoa(count))
   104  	}
   105  	return method
   106  }
   107  
   108  func GenName(prefix, suffix string) func(name string) string {
   109  	return func(name string) string {
   110  		return prefix + name + suffix
   111  	}
   112  }