Go语言 在gzip和csv阅读器之间动态切换

t1qtbnec  于 2023-01-28  发布在  Go
关注(0)|答案(2)|浏览(112)

我正在使用下面的代码片段来解压缩一个压缩的csv文件,并将其拆分为均匀的部分。

func split(feed MerchantFeed, db *gorm.DB) {
    var merchant Merchant
    db.First(&merchant, feed.MerchantID)
    merchant.Title = strings.Replace(merchant.Title, ".", "_", -1)
    fileID := merchant.Title + "__" + fmt.Sprint(feed.MerchantID) + "__" + uniuri.NewLen(20)
    out, _ := os.Create("/tmp/feeds/" + fileID)

    resp, err := http.Get(feed.Url)
    if err != nil {
        fmt.Println("Feed download error!")
    }
    defer resp.Body.Close()

    // Calling Copy method with its parameters
    bytes, err := io.Copy(out, resp.Body)

    if err != nil {
        fmt.Println("Could not create feed csv")
    }

    defer out.Close()

    in, _ := os.Open("/tmp/feeds/" + fileID)

    reader, err := gzip.NewReader(in)

    if err != nil {
        fmt.Println("Could not unzip feed " + fmt.Sprint(feed.MerchantID))
        err.Error()
        return
    }

    defer in.Close()

    // Empty byte slice.
    csvString := make([]byte, 100)

    // Read in data.
    _, errRead := reader.Read(csvString)
    if errRead != nil {
        fmt.Println(err)
    }
    path := "/tmp/split/feed__" + fileID + "__.csv"
    csvfile, err := os.Create(path)

    io.Copy(csvfile, reader)

    if fileID != "" {
        _, errx := splitter.Split(path, splitDir)

        if errx != nil {
            fmt.Println(errx)
        }

        defer csvfile.Close()
        //fmt.Println(result)

        result := os.Remove(path)
        if result != nil {
            fmt.Println(result)
        }
    }

    fmt.Printf("Size in MB (Feed %d): %d\n", feed.MerchantID, bytes)
    if bytes == 64 {
        fmt.Println("could not download feed from " + fmt.Sprint(feed.MerchantID))
        e := os.Remove("/tmp/feeds/" + fileID + ".gz")
        if e != nil {
            log.Fatal(e)
        }
        return
    }

}

到目前为止一切顺利。一切正常。但是有些商家把他们的csv打包成zip,有些则不是。所以我的想法是在gzip.NewReader()和csv.NewReader()之间动态切换,但是我不知道正确的方法。我想尽可能少地重复代码。所以我试着说

var zipNetworks []int
    zipNetworks = append(zipNetworks, 2)
    
    if slices.Contains(zipMerchants,merchant.Network_ID) {
        reader := zip.NewReader(in)
    } else {
        reader := csv.NewReader(in)
    }

但我的瞄准镜有问题。
有什么我不知道的聪明办法吗?

2sbarzqh

2sbarzqh1#

您可以使用通用的io.Reader接口。

var reader io.Reader
if slices.Contains(zipMerchants, merchant.Network_ID) {
    reader = zip.NewReader(in)
} else {
    reader = csv.NewReader(in)
}
wlzqhblo

wlzqhblo2#

将reader设置为文件。如果是gzip,则将reader替换为文件周围的gzip reader:

var reader io.Reader = in
if shouldUseGzip {
   reader, err = gzip.NewReader(reader)
   if err != nil {
       fmt.Println("Could not unzip feed " + fmt.Sprint(feed.MerchantID), err)
       return
}

相关问题