Connecting up the augmentor

Workiva · Jan 21, 2025 · df8064f · df8064f
1 parent 13c626e
commit df8064f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/build/augmentor.go b/build/augmentor.go
@@ -65,7 +65,7 @@ type Augmentor struct {
 	packages map[string]*pkgOverrideInfo
 }
 
-func (aug *Augmentor) Augment(xctx XContext, pkg *PackageData, fileSet *token.FileSet, file *ast.File) {
+func (aug *Augmentor) Augment(xctx XContext, pkg *PackageData, fileSet *token.FileSet, file *ast.File) error {
 	pkgAug := aug.getPackageOverrides(xctx, pkg, fileSet)
 
 	augmentOriginalImports(pkg.ImportPath, file)
@@ -80,10 +80,20 @@ func (aug *Augmentor) Augment(xctx XContext, pkg *PackageData, fileSet *token.Fi
 		// additional methods and information from the natives.
 		err := astutil.ConcatenateFiles(file, pkgAug.overlayFiles...)
 		if err != nil {
-			panic(fmt.Errorf("failed to concatenate overlay files onto %q: %w", fileSet.Position(file.Package).Filename, err))
+			return fmt.Errorf("failed to concatenate overlay files onto %q: %w", fileSet.Position(file.Package).Filename, err)
 		}
 		pkgAug.overlayFiles = nil
 	}
+
+	return nil
+}
+
+func (aug *Augmentor) GetJSFiles(pkg *PackageData) []JSFile {
+	pkgAug, ok := aug.packages[pkg.ImportPath]
+	if !ok {
+		return nil
+	}
+	return pkgAug.jsFiles
 }
 
 // getPackageOverrides looks up an already loaded package override

diff --git a/build/build.go b/build/build.go
@@ -151,10 +151,13 @@ func parseAndAugment(xctx XContext, pkg *PackageData, fileSet *token.FileSet) ([
 
 	aug := &Augmentor{}
 	for _, file := range originalFiles {
-		aug.Augment(xctx, pkg, fileSet, file)
+		err := aug.Augment(xctx, pkg, fileSet, file)
+		if err != nil {
+			return nil, nil, err
+		}
 	}
 
-	jsFiles := aug.packages[pkg.ImportPath].jsFiles
+	jsFiles := aug.GetJSFiles(pkg)
 	return originalFiles, jsFiles, nil
 }