@@ -166,6 +166,14 @@ func AppendFuncAtLine(code string, lineNumber uint64) FunctionOptions {
166166
167167// AppendInsideFuncCall adds an argument to a function call. For instances, the method have a parameter a
168168// // call 'New(param1, param2)' and we want to add the param3 the result will be 'New(param1, param2, param3)'.
169+ // AppendInsideFuncCall appends code inside a function call.
170+ // The callName parameter can be either:
171+ // - Simple name: "NewKeeper" matches any call to NewKeeper regardless of package/receiver
172+ // - Qualified name: "foo.NewKeeper" matches only calls to NewKeeper with foo as the package/receiver
173+ //
174+ // The code parameter is the argument to insert, and index specifies the position:
175+ // - index >= 0: insert at the specified position
176+ // - index == -1: append at the end
169177func AppendInsideFuncCall (callName , code string , index int ) FunctionOptions {
170178 return func (c * functionOpts ) {
171179 c .insideCall = append (c .insideCall , functionCall {
@@ -177,8 +185,12 @@ func AppendInsideFuncCall(callName, code string, index int) FunctionOptions {
177185}
178186
179187// AppendFuncStruct adds a field to a struct literal. For instance,
180- // // the struct has only one parameter 'Params{Param1: param1}' and we want to add
181- // // the param2 the result will be 'Params{Param1: param1, Param2: param2}'.
188+ // the struct has only one parameter 'Params{Param1: param1}' and we want to add
189+ // the param2 the result will be 'Params{Param1: param1, Param2: param2}'.
190+ //
191+ // The name parameter can be either:
192+ // - Simple name: "Keeper" matches any struct literal of type Keeper regardless of package
193+ // - Qualified name: "keeper.Keeper" matches only struct literals with keeper as the package
182194func AppendFuncStruct (name , param , code string ) FunctionOptions {
183195 return func (c * functionOpts ) {
184196 c .insideStruct = append (c .insideStruct , functionStruct {
@@ -233,7 +245,7 @@ func ModifyFunction(content string, funcName string, functions ...FunctionOption
233245 fset := token .NewFileSet ()
234246 file , err := parser .ParseFile (fset , "" , content , parser .ParseComments )
235247 if err != nil {
236- return "" , errors .Errorf ("failed to parse file: %w" , err )
248+ return "" , errors .Errorf ("failed to parse file (%s) : %w" , funcName , err )
237249 }
238250
239251 cmap := ast .NewCommentMap (fset , file , file .Comments )
@@ -363,7 +375,7 @@ func addNewLine(fileSet *token.FileSet, funcDecl *ast.FuncDecl, newLines []funct
363375 for _ , newLine := range newLines {
364376 // Validate line number
365377 if newLine .number > uint64 (len (funcDecl .Body .List ))- 1 {
366- return errors .Errorf ("line number %d out of range" , newLine .number )
378+ return errors .Errorf ("line number %d out of range (max %d) " , newLine .number , len ( funcDecl . Body . List ) - 1 )
367379 }
368380
369381 // Parse insertion code
@@ -484,6 +496,12 @@ func exprName(expr ast.Expr) (string, bool) {
484496 case * ast.Ident :
485497 return exp .Name , true
486498 case * ast.SelectorExpr :
499+ // Check if X is an identifier to get the package name
500+ if ident , ok := exp .X .(* ast.Ident ); ok {
501+ // Return qualified name: package.Function
502+ return ident .Name + "." + exp .Sel .Name , true
503+ }
504+ // Fallback to just the selector name if X is not an identifier
487505 return exp .Sel .Name , true
488506 default :
489507 return "" , false
@@ -691,30 +709,58 @@ func applyFunctionOptions(fileSet *token.FileSet, f *ast.FuncDecl, opts *functio
691709 return true
692710 }
693711
694- calls , ok := callMap [name ]
695- if ! ok {
712+ // Collect all matching calls (both qualified and unqualified names)
713+ var allCalls functionCalls
714+ if calls , ok := callMap [name ]; ok {
715+ allCalls = append (allCalls , calls ... )
716+ delete (callMapCheck , name )
717+ }
718+
719+ // Also check for unqualified name if this is a selector expression
720+ if sel , isSel := expr .Fun .(* ast.SelectorExpr ); isSel {
721+ simpleName := sel .Sel .Name
722+ if calls , ok := callMap [simpleName ]; ok {
723+ allCalls = append (allCalls , calls ... )
724+ delete (callMapCheck , simpleName )
725+ }
726+ }
727+
728+ if len (allCalls ) == 0 {
696729 return true
697730 }
698731
699- if err := addFunctionCall (expr , calls ); err != nil {
732+ if err := addFunctionCall (expr , allCalls ); err != nil {
700733 errInspect = err
701734 return false
702735 }
703- delete (callMapCheck , name )
704736
705737 case * ast.CompositeLit :
706738 name , exist := exprName (expr .Type )
707739 if ! exist {
708740 return true
709741 }
710742
711- structs , ok := structMap [name ]
712- if ! ok {
743+ // Collect all matching structs (both qualified and unqualified names)
744+ var allStructs functionStructs
745+ if structs , ok := structMap [name ]; ok {
746+ allStructs = append (allStructs , structs ... )
747+ delete (structMapCheck , name )
748+ }
749+
750+ // Also check for unqualified name if this is a selector expression
751+ if sel , isSel := expr .Type .(* ast.SelectorExpr ); isSel {
752+ simpleName := sel .Sel .Name
753+ if structs , ok := structMap [simpleName ]; ok {
754+ allStructs = append (allStructs , structs ... )
755+ delete (structMapCheck , simpleName )
756+ }
757+ }
758+
759+ if len (allStructs ) == 0 {
713760 return true
714761 }
715762
716- addStructs (fileSet , expr , structs )
717- delete (structMapCheck , name )
763+ addStructs (fileSet , expr , allStructs )
718764
719765 default :
720766 return true
0 commit comments