ioutil.go

Documentation: github.com/spf13/afero

     1  // Copyright ©2015 The Go Authors
     2  // Copyright ©2015 Steve Francia <spf@spf13.com>
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package afero
    17  
    18  import (
    19  	"bytes"
    20  	"io"
    21  	"os"
    22  	"path/filepath"
    23  	"sort"
    24  	"strconv"
    25  	"strings"
    26  	"sync"
    27  	"time"
    28  )
    29  
    30  // byName implements sort.Interface.
    31  type byName []os.FileInfo
    32  
    33  func (f byName) Len() int           { return len(f) }
    34  func (f byName) Less(i, j int) bool { return f[i].Name() < f[j].Name() }
    35  func (f byName) Swap(i, j int)      { f[i], f[j] = f[j], f[i] }
    36  
    37  // ReadDir reads the directory named by dirname and returns
    38  // a list of sorted directory entries.
    39  func (a Afero) ReadDir(dirname string) ([]os.FileInfo, error) {
    40  	return ReadDir(a.Fs, dirname)
    41  }
    42  
    43  func ReadDir(fs Fs, dirname string) ([]os.FileInfo, error) {
    44  	f, err := fs.Open(dirname)
    45  	if err != nil {
    46  		return nil, err
    47  	}
    48  	list, err := f.Readdir(-1)
    49  	f.Close()
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	sort.Sort(byName(list))
    54  	return list, nil
    55  }
    56  
    57  // ReadFile reads the file named by filename and returns the contents.
    58  // A successful call returns err == nil, not err == EOF. Because ReadFile
    59  // reads the whole file, it does not treat an EOF from Read as an error
    60  // to be reported.
    61  func (a Afero) ReadFile(filename string) ([]byte, error) {
    62  	return ReadFile(a.Fs, filename)
    63  }
    64  
    65  func ReadFile(fs Fs, filename string) ([]byte, error) {
    66  	f, err := fs.Open(filename)
    67  	if err != nil {
    68  		return nil, err
    69  	}
    70  	defer f.Close()
    71  	// It's a good but not certain bet that FileInfo will tell us exactly how much to
    72  	// read, so let's try it but be prepared for the answer to be wrong.
    73  	var n int64
    74  
    75  	if fi, err := f.Stat(); err == nil {
    76  		// Don't preallocate a huge buffer, just in case.
    77  		if size := fi.Size(); size < 1e9 {
    78  			n = size
    79  		}
    80  	}
    81  	// As initial capacity for readAll, use n + a little extra in case Size is zero,
    82  	// and to avoid another allocation after Read has filled the buffer.  The readAll
    83  	// call will read into its allocated internal buffer cheaply.  If the size was
    84  	// wrong, we'll either waste some space off the end or reallocate as needed, but
    85  	// in the overwhelmingly common case we'll get it just right.
    86  	return readAll(f, n+bytes.MinRead)
    87  }
    88  
    89  // readAll reads from r until an error or EOF and returns the data it read
    90  // from the internal buffer allocated with a specified capacity.
    91  func readAll(r io.Reader, capacity int64) (b []byte, err error) {
    92  	buf := bytes.NewBuffer(make([]byte, 0, capacity))
    93  	// If the buffer overflows, we will get bytes.ErrTooLarge.
    94  	// Return that as an error. Any other panic remains.
    95  	defer func() {
    96  		e := recover()
    97  		if e == nil {
    98  			return
    99  		}
   100  		if panicErr, ok := e.(error); ok && panicErr == bytes.ErrTooLarge {
   101  			err = panicErr
   102  		} else {
   103  			panic(e)
   104  		}
   105  	}()
   106  	_, err = buf.ReadFrom(r)
   107  	return buf.Bytes(), err
   108  }
   109  
   110  // ReadAll reads from r until an error or EOF and returns the data it read.
   111  // A successful call returns err == nil, not err == EOF. Because ReadAll is
   112  // defined to read from src until EOF, it does not treat an EOF from Read
   113  // as an error to be reported.
   114  func ReadAll(r io.Reader) ([]byte, error) {
   115  	return readAll(r, bytes.MinRead)
   116  }
   117  
   118  // WriteFile writes data to a file named by filename.
   119  // If the file does not exist, WriteFile creates it with permissions perm;
   120  // otherwise WriteFile truncates it before writing.
   121  func (a Afero) WriteFile(filename string, data []byte, perm os.FileMode) error {
   122  	return WriteFile(a.Fs, filename, data, perm)
   123  }
   124  
   125  func WriteFile(fs Fs, filename string, data []byte, perm os.FileMode) error {
   126  	f, err := fs.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
   127  	if err != nil {
   128  		return err
   129  	}
   130  	n, err := f.Write(data)
   131  	if err == nil && n < len(data) {
   132  		err = io.ErrShortWrite
   133  	}
   134  	if err1 := f.Close(); err == nil {
   135  		err = err1
   136  	}
   137  	return err
   138  }
   139  
   140  // Random number state.
   141  // We generate random temporary file names so that there's a good
   142  // chance the file doesn't exist yet - keeps the number of tries in
   143  // TempFile to a minimum.
   144  var (
   145  	randNum uint32
   146  	randmu  sync.Mutex
   147  )
   148  
   149  func reseed() uint32 {
   150  	return uint32(time.Now().UnixNano() + int64(os.Getpid()))
   151  }
   152  
   153  func nextRandom() string {
   154  	randmu.Lock()
   155  	r := randNum
   156  	if r == 0 {
   157  		r = reseed()
   158  	}
   159  	r = r*1664525 + 1013904223 // constants from Numerical Recipes
   160  	randNum = r
   161  	randmu.Unlock()
   162  	return strconv.Itoa(int(1e9 + r%1e9))[1:]
   163  }
   164  
   165  // TempFile creates a new temporary file in the directory dir,
   166  // opens the file for reading and writing, and returns the resulting *os.File.
   167  // The filename is generated by taking pattern and adding a random
   168  // string to the end. If pattern includes a "*", the random string
   169  // replaces the last "*".
   170  // If dir is the empty string, TempFile uses the default directory
   171  // for temporary files (see os.TempDir).
   172  // Multiple programs calling TempFile simultaneously
   173  // will not choose the same file. The caller can use f.Name()
   174  // to find the pathname of the file. It is the caller's responsibility
   175  // to remove the file when no longer needed.
   176  func (a Afero) TempFile(dir, pattern string) (f File, err error) {
   177  	return TempFile(a.Fs, dir, pattern)
   178  }
   179  
   180  func TempFile(fs Fs, dir, pattern string) (f File, err error) {
   181  	if dir == "" {
   182  		dir = os.TempDir()
   183  	}
   184  
   185  	var prefix, suffix string
   186  	if pos := strings.LastIndex(pattern, "*"); pos != -1 {
   187  		prefix, suffix = pattern[:pos], pattern[pos+1:]
   188  	} else {
   189  		prefix = pattern
   190  	}
   191  
   192  	nconflict := 0
   193  	for i := 0; i < 10000; i++ {
   194  		name := filepath.Join(dir, prefix+nextRandom()+suffix)
   195  		f, err = fs.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600)
   196  		if os.IsExist(err) {
   197  			if nconflict++; nconflict > 10 {
   198  				randmu.Lock()
   199  				randNum = reseed()
   200  				randmu.Unlock()
   201  			}
   202  			continue
   203  		}
   204  		break
   205  	}
   206  	return
   207  }
   208  
   209  // TempDir creates a new temporary directory in the directory dir
   210  // with a name beginning with prefix and returns the path of the
   211  // new directory.  If dir is the empty string, TempDir uses the
   212  // default directory for temporary files (see os.TempDir).
   213  // Multiple programs calling TempDir simultaneously
   214  // will not choose the same directory.  It is the caller's responsibility
   215  // to remove the directory when no longer needed.
   216  func (a Afero) TempDir(dir, prefix string) (name string, err error) {
   217  	return TempDir(a.Fs, dir, prefix)
   218  }
   219  
   220  func TempDir(fs Fs, dir, prefix string) (name string, err error) {
   221  	if dir == "" {
   222  		dir = os.TempDir()
   223  	}
   224  
   225  	nconflict := 0
   226  	for i := 0; i < 10000; i++ {
   227  		try := filepath.Join(dir, prefix+nextRandom())
   228  		err = fs.Mkdir(try, 0o700)
   229  		if os.IsExist(err) {
   230  			if nconflict++; nconflict > 10 {
   231  				randmu.Lock()
   232  				randNum = reseed()
   233  				randmu.Unlock()
   234  			}
   235  			continue
   236  		}
   237  		if err == nil {
   238  			name = try
   239  		}
   240  		break
   241  	}
   242  	return
   243  }
   244
View as plain text