ioutil.go 7.0 KB

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