本站源代码
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

171 lines
4.5KB

  1. // Package ioutil implements some I/O utility functions.
  2. package ioutil
  3. import (
  4. "bufio"
  5. "context"
  6. "errors"
  7. "io"
  8. "github.com/jbenet/go-context/io"
  9. )
  10. type readPeeker interface {
  11. io.Reader
  12. Peek(int) ([]byte, error)
  13. }
  14. var (
  15. ErrEmptyReader = errors.New("reader is empty")
  16. )
  17. // NonEmptyReader takes a reader and returns it if it is not empty, or
  18. // `ErrEmptyReader` if it is empty. If there is an error when reading the first
  19. // byte of the given reader, it will be propagated.
  20. func NonEmptyReader(r io.Reader) (io.Reader, error) {
  21. pr, ok := r.(readPeeker)
  22. if !ok {
  23. pr = bufio.NewReader(r)
  24. }
  25. _, err := pr.Peek(1)
  26. if err == io.EOF {
  27. return nil, ErrEmptyReader
  28. }
  29. if err != nil {
  30. return nil, err
  31. }
  32. return pr, nil
  33. }
  34. type readCloser struct {
  35. io.Reader
  36. closer io.Closer
  37. }
  38. func (r *readCloser) Close() error {
  39. return r.closer.Close()
  40. }
  41. // NewReadCloser creates an `io.ReadCloser` with the given `io.Reader` and
  42. // `io.Closer`.
  43. func NewReadCloser(r io.Reader, c io.Closer) io.ReadCloser {
  44. return &readCloser{Reader: r, closer: c}
  45. }
  46. type writeCloser struct {
  47. io.Writer
  48. closer io.Closer
  49. }
  50. func (r *writeCloser) Close() error {
  51. return r.closer.Close()
  52. }
  53. // NewWriteCloser creates an `io.WriteCloser` with the given `io.Writer` and
  54. // `io.Closer`.
  55. func NewWriteCloser(w io.Writer, c io.Closer) io.WriteCloser {
  56. return &writeCloser{Writer: w, closer: c}
  57. }
  58. type writeNopCloser struct {
  59. io.Writer
  60. }
  61. func (writeNopCloser) Close() error { return nil }
  62. // WriteNopCloser returns a WriteCloser with a no-op Close method wrapping
  63. // the provided Writer w.
  64. func WriteNopCloser(w io.Writer) io.WriteCloser {
  65. return writeNopCloser{w}
  66. }
  67. // CheckClose calls Close on the given io.Closer. If the given *error points to
  68. // nil, it will be assigned the error returned by Close. Otherwise, any error
  69. // returned by Close will be ignored. CheckClose is usually called with defer.
  70. func CheckClose(c io.Closer, err *error) {
  71. if cerr := c.Close(); cerr != nil && *err == nil {
  72. *err = cerr
  73. }
  74. }
  75. // NewContextWriter wraps a writer to make it respect given Context.
  76. // If there is a blocking write, the returned Writer will return whenever the
  77. // context is cancelled (the return values are n=0 and err=ctx.Err()).
  78. func NewContextWriter(ctx context.Context, w io.Writer) io.Writer {
  79. return ctxio.NewWriter(ctx, w)
  80. }
  81. // NewContextReader wraps a reader to make it respect given Context.
  82. // If there is a blocking read, the returned Reader will return whenever the
  83. // context is cancelled (the return values are n=0 and err=ctx.Err()).
  84. func NewContextReader(ctx context.Context, r io.Reader) io.Reader {
  85. return ctxio.NewReader(ctx, r)
  86. }
  87. // NewContextWriteCloser as NewContextWriter but with io.Closer interface.
  88. func NewContextWriteCloser(ctx context.Context, w io.WriteCloser) io.WriteCloser {
  89. ctxw := ctxio.NewWriter(ctx, w)
  90. return NewWriteCloser(ctxw, w)
  91. }
  92. // NewContextReadCloser as NewContextReader but with io.Closer interface.
  93. func NewContextReadCloser(ctx context.Context, r io.ReadCloser) io.ReadCloser {
  94. ctxr := ctxio.NewReader(ctx, r)
  95. return NewReadCloser(ctxr, r)
  96. }
  97. type readerOnError struct {
  98. io.Reader
  99. notify func(error)
  100. }
  101. // NewReaderOnError returns a io.Reader that call the notify function when an
  102. // unexpected (!io.EOF) error happens, after call Read function.
  103. func NewReaderOnError(r io.Reader, notify func(error)) io.Reader {
  104. return &readerOnError{r, notify}
  105. }
  106. // NewReadCloserOnError returns a io.ReadCloser that call the notify function
  107. // when an unexpected (!io.EOF) error happens, after call Read function.
  108. func NewReadCloserOnError(r io.ReadCloser, notify func(error)) io.ReadCloser {
  109. return NewReadCloser(NewReaderOnError(r, notify), r)
  110. }
  111. func (r *readerOnError) Read(buf []byte) (n int, err error) {
  112. n, err = r.Reader.Read(buf)
  113. if err != nil && err != io.EOF {
  114. r.notify(err)
  115. }
  116. return
  117. }
  118. type writerOnError struct {
  119. io.Writer
  120. notify func(error)
  121. }
  122. // NewWriterOnError returns a io.Writer that call the notify function when an
  123. // unexpected (!io.EOF) error happens, after call Write function.
  124. func NewWriterOnError(w io.Writer, notify func(error)) io.Writer {
  125. return &writerOnError{w, notify}
  126. }
  127. // NewWriteCloserOnError returns a io.WriteCloser that call the notify function
  128. //when an unexpected (!io.EOF) error happens, after call Write function.
  129. func NewWriteCloserOnError(w io.WriteCloser, notify func(error)) io.WriteCloser {
  130. return NewWriteCloser(NewWriterOnError(w, notify), w)
  131. }
  132. func (r *writerOnError) Write(p []byte) (n int, err error) {
  133. n, err = r.Writer.Write(p)
  134. if err != nil && err != io.EOF {
  135. r.notify(err)
  136. }
  137. return
  138. }
上海开阖软件有限公司 沪ICP备12045867号-1