github.com/replit/upm@v0.0.0-20240423230255-9ce4fc3ea24c/internal/backends/python/gen_pypi_map/archive_reader.go (about)

     1  package main
     2  
     3  import (
     4  	"archive/tar"
     5  	"archive/zip"
     6  	"compress/gzip"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  	"os/exec"
    11  )
    12  
    13  type ArchiveReader struct {
    14  	Next   func() bool
    15  	Error  func() error
    16  	File   func() string
    17  	Reader func() (io.Reader, error)
    18  	Close  func()
    19  	Dump   func(dir string) error
    20  }
    21  
    22  func MakeZipReader(reader io.Reader, size int64) (ArchiveReader, error) {
    23  	cacheFile, err := os.CreateTemp(os.TempDir(), "pypi-")
    24  	if err != nil {
    25  		return ArchiveReader{}, fmt.Errorf("Cannot open cache for zip: %v", err)
    26  	}
    27  
    28  	// Copy the contents of the reader to disk
    29  	_, err = io.Copy(cacheFile, reader)
    30  	if err != nil {
    31  		return ArchiveReader{}, err
    32  	}
    33  
    34  	cacheFile.Close()
    35  
    36  	zipFile, err := os.Open(cacheFile.Name())
    37  	if err != nil {
    38  		return ArchiveReader{}, fmt.Errorf("Cannot open cache for zip: %v", err)
    39  	}
    40  
    41  	zipReader, err := zip.NewReader(zipFile, size)
    42  	if err != nil {
    43  		return ArchiveReader{}, err
    44  	}
    45  
    46  	i := -1
    47  	next := func() bool {
    48  		i++
    49  		return i < len(zipReader.File)
    50  	}
    51  
    52  	archiveFile := func() string {
    53  		return zipReader.File[i].Name
    54  	}
    55  
    56  	archiveError := func() error {
    57  		if i >= len(zipReader.File) {
    58  			return io.EOF
    59  		}
    60  		return nil
    61  	}
    62  
    63  	fileReader := func() (io.Reader, error) {
    64  		return zipReader.File[i].Open()
    65  	}
    66  
    67  	archiveClose := func() {
    68  		zipFile.Close()
    69  		os.Remove(zipFile.Name())
    70  	}
    71  
    72  	archiveDump := func(dir string) error {
    73  		cmd := exec.Command("unzip", "-d", dir, cacheFile.Name())
    74  		return cmd.Run()
    75  	}
    76  
    77  	return ArchiveReader{Next: next, Error: archiveError, File: archiveFile, Reader: fileReader, Close: archiveClose, Dump: archiveDump}, nil
    78  }
    79  
    80  func MakeTarballReader(reader io.Reader) (ArchiveReader, error) {
    81  	gzipReader, err := gzip.NewReader(reader)
    82  	if err != nil {
    83  		return ArchiveReader{}, err
    84  	}
    85  
    86  	tarReader := tar.NewReader(gzipReader)
    87  	var header *tar.Header
    88  
    89  	next := func() bool {
    90  		header, err = tarReader.Next()
    91  		return err != nil
    92  	}
    93  
    94  	tarballError := func() error {
    95  		return err
    96  	}
    97  
    98  	tarballFile := func() string {
    99  		return header.Name
   100  	}
   101  
   102  	fileReader := func() (io.Reader, error) {
   103  		return tarReader, nil
   104  	}
   105  
   106  	closeGzip := func() {
   107  		gzipReader.Close()
   108  	}
   109  
   110  	archiveDump := func(dir string) error {
   111  		cmd := exec.Command("tar", "-x", "-C", dir)
   112  		p, err := cmd.StdinPipe()
   113  		if err != nil {
   114  			return err
   115  		}
   116  
   117  		defer p.Close()
   118  
   119  		err = cmd.Start()
   120  		if err != nil {
   121  			return err
   122  		}
   123  
   124  		_, err = io.Copy(p, gzipReader)
   125  		if err != nil {
   126  			return err
   127  		}
   128  
   129  		err = cmd.Wait()
   130  		if err != nil {
   131  			return err
   132  		}
   133  
   134  		return nil
   135  	}
   136  
   137  	return ArchiveReader{Next: next, Error: tarballError, File: tarballFile, Reader: fileReader, Close: closeGzip, Dump: archiveDump}, nil
   138  }