From 09d81f8553e6725d7e3cf9d20dc23bbcff51da36 Mon Sep 17 00:00:00 2001 From: Devang Gaur Date: Thu, 12 Nov 2020 14:18:24 +0530 Subject: [PATCH] optimized yaml utility function --- pkg/utils/yaml.go | 154 +++++++++++++++------------------------------- 1 file changed, 49 insertions(+), 105 deletions(-) diff --git a/pkg/utils/yaml.go b/pkg/utils/yaml.go index 2fb8238d4..c208d81e8 100644 --- a/pkg/utils/yaml.go +++ b/pkg/utils/yaml.go @@ -26,118 +26,76 @@ var ( func LoadYAML(filePath string) ([]*IacDocument, error) { iacDocumentList := make([]*IacDocument, 0) - // First pass determines line number data - { // Limit the scope for Close() - file, err := os.Open(filePath) - if err != nil { - return iacDocumentList, err - } - defer file.Close() - - scanner := bufio.NewScanner(file) - startLineNumber := 1 - currentLineNumber := 1 - for scanner.Scan() { - if strings.HasPrefix(scanner.Text(), "---") { - // We've found the end-of-directives marker, so record results for the current document - iacDocumentList = append(iacDocumentList, &IacDocument{ - Type: YAMLDoc, - StartLine: startLineNumber, - EndLine: currentLineNumber, - FilePath: filePath, - }) - startLineNumber = currentLineNumber + 1 - } - currentLineNumber++ - } - - // Add the very last entry - iacDocumentList = append(iacDocumentList, &IacDocument{ - Type: YAMLDoc, - StartLine: startLineNumber, - EndLine: currentLineNumber, - FilePath: filePath, - }) - - if err = scanner.Err(); err != nil { - return iacDocumentList, err - } + file, err := os.Open(filePath) + if err != nil { + return iacDocumentList, err } + defer file.Close() - // Second pass extracts all YAML documents and saves it in the document struct fileBytes, err := ioutil.ReadFile(filePath) if err != nil { return iacDocumentList, err } - dec := yaml.NewDecoder(bytes.NewReader(fileBytes)) - i := 0 - for { - // each iteration extracts and marshals one yaml document - var value interface{} - err = dec.Decode(&value) - if err == io.EOF { - break - } - if err != nil { - return iacDocumentList, err - } - if i > (len(iacDocumentList) - 1) { - return iacDocumentList, errHighDocumentCount - } - - var documentBytes []byte - documentBytes, err = yaml.Marshal(value) - if err != nil { - return iacDocumentList, err - } - iacDocumentList[i].Data = documentBytes - i++ - } - - return iacDocumentList, nil + return getIacDocumentList(bufio.NewScanner(file), fileBytes, filePath) } // LoadYAMLString loads a YAML String. Can return one or more IaC Documents. // Besides reading in file data, its main purpose is to determine and store line number and filename metadata func LoadYAMLString(data, absFilePath string) ([]*IacDocument, error) { + return getIacDocumentList(bufio.NewScanner(strings.NewReader(data)), []byte(data), absFilePath) +} + +// ReadYamlFile reads a yaml file and load content in a map[string]interface{} type +func ReadYamlFile(path string) (map[string]interface{}, error) { + dat, err := ioutil.ReadFile(path) + if err != nil { + return nil, err + } + var output map[string]interface{} + err = yaml.Unmarshal(dat, &output) + if err != nil { + return nil, err + } + return output, nil +} + +// getIacDocumentList provides one or more IaC Documents. +// Besides reading in file data, its main purpose is to determine and store line number and filename metadata +func getIacDocumentList(scanner *bufio.Scanner, bytearray []byte, filePath string) ([]*IacDocument, error) { iacDocumentList := make([]*IacDocument, 0) // First pass determines line number data - { // Limit the scope for Close() - scanner := bufio.NewScanner(strings.NewReader(data)) - startLineNumber := 1 - currentLineNumber := 1 - for scanner.Scan() { - if strings.HasPrefix(scanner.Text(), "---") { - // We've found the end-of-directives marker, so record results for the current document - iacDocumentList = append(iacDocumentList, &IacDocument{ - Type: YAMLDoc, - StartLine: startLineNumber, - EndLine: currentLineNumber, - FilePath: absFilePath, - }) - startLineNumber = currentLineNumber + 1 - } - currentLineNumber++ + startLineNumber := 1 + currentLineNumber := 1 + for scanner.Scan() { + if strings.HasPrefix(scanner.Text(), "---") { + // We've found the end-of-directives marker, so record results for the current document + iacDocumentList = append(iacDocumentList, &IacDocument{ + Type: YAMLDoc, + StartLine: startLineNumber, + EndLine: currentLineNumber, + FilePath: filePath, + }) + startLineNumber = currentLineNumber + 1 } + currentLineNumber++ + } - // Add the very last entry - iacDocumentList = append(iacDocumentList, &IacDocument{ - Type: YAMLDoc, - StartLine: startLineNumber, - EndLine: currentLineNumber, - FilePath: "", - }) + // Add the very last entry + iacDocumentList = append(iacDocumentList, &IacDocument{ + Type: YAMLDoc, + StartLine: startLineNumber, + EndLine: currentLineNumber, + FilePath: "", + }) - if err := scanner.Err(); err != nil { - return iacDocumentList, err - } + if err := scanner.Err(); err != nil { + return iacDocumentList, err } - // Second pass extracts all YAML documents and saves it in the document struct + dec := yaml.NewDecoder(bytes.NewReader(bytearray)) - dec := yaml.NewDecoder(bytes.NewReader([]byte(data))) i := 0 for { // each iteration extracts and marshals one yaml document @@ -164,17 +122,3 @@ func LoadYAMLString(data, absFilePath string) ([]*IacDocument, error) { return iacDocumentList, nil } - -// ReadYamlFile reads a yaml file and load content in a map[string]interface{} type -func ReadYamlFile(path string) (map[string]interface{}, error) { - dat, err := ioutil.ReadFile(path) - if err != nil { - return nil, err - } - var output map[string]interface{} - err = yaml.Unmarshal(dat, &output) - if err != nil { - return nil, err - } - return output, nil -}