<commit_before>//
// s3 - Swiss army pen-knife for Amazon S3.
//
//   https://github.com/barnybug/s3
//
// Copyright (c) 2015 <NAME>

package s3

import (
	"bytes"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

var reBucketPath = regexp.MustCompile("^(?:s3://)?([^/]+)/?(.*)$")
var out io.Writer = os.Stdout
var err io.Writer = os.Stderr

var (
	ErrNotFound = errors.New("No files found")
)

func extractBucketPath(url string) (string, string) {
	parts := reBucketPath.FindStringSubmatch(url)
	return parts[1], parts[2]
}

func listBuckets(conn s3iface.S3API) error {
	output, err := conn.ListBuckets(nil)
	if err != nil {
		return err
	}
	for _, b := range output.Buckets {
		fmt.Fprintf(out, "s3://%s/\n", *b.Name)
	}
	return nil
}

func iterateKeys(conn s3iface.S3API, urls []string, callback func(file File) error) error {
	found := false
	for _, url := range urls {
		fs := getFilesystem(conn, url)
		ch := fs.Files()
		for file := range ch {
			found = true
			err := callback(file)
			if err != nil {
				return err
			}
		}
		if fs.Error() != nil {
			return fs.Error()
		}
	}
	if !found {
		return ErrNotFound
	}
	return nil
}

func iterateKeysParallel(conn s3iface.S3API, urls []string, callback func(file File) error) error {
	// create pool for processing
	var err error
	wg := sync.WaitGroup{}
	q := make(chan File, 1000)
	for i := 0; i < parallel; i += 1 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for key := range q {
				e := callback(key)
				if e != nil {
					err = e
					return
				}
			}
		}()
	}

	e := iterateKeys(conn, urls, func(file File) error {
		q <- file
		return nil
	})
	if e != nil {
		return e
	}

	close(q)
	wg.Wait()
	return err
}

func listKeys(conn s3iface.S3API, urls []string) error {
	var count, totalSize int64
	err := iterateKeys(conn, urls, func(file File) error {
		if quiet {
			fmt.Fprintln(out, file)
		} else {
			fmt.Fprintf(out, "%s\t%db\n", file, file.Size())
		}
		count += 1
		totalSize += file.Size()
		return nil
	})
	if err != nil && err != ErrNotFound {
		return err
	}
	if !quiet {
		fmt.Fprintf(out, "\n%d files, %d bytes\n", count, totalSize)
	}
	return nil
}

func getKeys(conn s3iface.S3API, urls []string) error {
	for _, url := range urls {
		if !isS3Url(url) {
			return errors.New("s3:// url required")
		}
	}

	err := iterateKeysParallel(conn, urls, func(file File) error {
		reader, err := file.Reader()
		if err != nil {
			return err
		}
		defer reader.Close()

		// write files under relative path to the source path
		fpath := file.Relative()
		dirpath := path.Dir(fpath)
		if dirpath != "." {
			err = os.MkdirAll(dirpath, 0777)
			if err != nil {
				return err
			}
		}

		writer, err := os.Create(fpath)
		if err != nil {
			return err
		}
		nbytes, err := io.Copy(writer, reader)
		if err != nil {
			return err
		}
		if !quiet {
			fmt.Fprintf(out, "%s -> %s (%d bytes)\n", file, fpath, nbytes)
		}
		return nil
	})
	return err
}

func catKeys(conn s3iface.S3API, urls []string) error {
	return iterateKeysParallel(conn, urls, func(file File) error {
		reader, err := file.Reader()
		if err != nil {
			return err
		}
		defer reader.Close()

		if strings.HasSuffix(file.String(), ".gz") {
			reader, err = gzip.NewReader(reader)
			if err != nil {
				return err
			}
		}

		_, err = io.Copy(out, reader)
		if err != nil {
			return err
		}
		return nil
	})
}

func outputMatches(buf []byte, needle []byte, prefix string) {
	p := 0
	for {
		i := bytes.Index(buf[p:], needle)
		if i == -1 {
			break
		}
		i += p
		lineStart := bytes.LastIndexByte(buf[:i], byte('\n'))
		if lineStart == -1 {
			lineStart = 0
		} else {
			lineStart += 1
		}
		lineEnd := bytes.IndexByte(buf[i:], byte('\n'))
		if lineEnd == -1 {
			lineEnd = len(buf)
		} else {
			lineEnd += i
		}
		line := string(buf[lineStart:lineEnd])

		fmt.Fprintf(out, "%s%s\n", prefix, line)

		p = lineEnd + 1
		if p > len(buf)-len(needle) {
			break
		}
	}
}

func grepKeys(conn s3iface.S3API, find string, urls []string, noKeysPrefix bool, keysWithMatches bool) error {
	needle := []byte(find)

	return iterateKeysParallel(conn, urls, func(file File) error {
		reader, err := file.Reader()
		if err != nil {
			return err
		}
		defer reader.Close()

		if strings.HasSuffix(file.String(), ".gz") {
			reader, err = gzip.NewReader(reader)
			if err != nil {
				return err
			}
		}

		prefix := ""
		if !noKeysPrefix {
			prefix = file.String() + ":"
		}

		buf := make([]byte, 4096)
		offset := 0
		for n, err := reader.Read(buf[offset:]); n > 0 && err == nil; n, err = reader.Read(buf) {
			if bytes.Contains(buf, needle) {
				if keysWithMatches {
					// only filename required, bail early
					fmt.Fprintln(out, file.String())
					break
				} else {
					outputMatches(buf, needle, prefix)
				}
			}
			// handle overlapping matches - copy last N-1 bytes to start of next
			offset = len(needle) - 1
			if offset > n {
				offset = n
			} else {
				copy(buf, buf[n-offset:])
			}
		}<commit_msg>Fix test<commit_after>//
// s3 - Swiss army pen-knife for Amazon S3.
//
//   https://github.com/barnybug/s3
//
// Copyright (c) 2015 <NAME>

package s3

import (
	"bytes"
	"compress/gzip"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3iface"
)

var reBucketPath = regexp.MustCompile("^(?:s3://)?([^/]+)/?(.*)$")
var out io.Writer = os.Stdout
var err io.Writer = os.Stderr

var (
	ErrNotFound = errors.New("No files found")
)

func extractBucketPath(url string) (string, string) {
	parts := reBucketPath.FindStringSubmatch(url)
	return parts[1], parts[2]
}

func listBuckets(conn s3iface.S3API) error {
	output, err := conn.ListBuckets(nil)
	if err != nil {
		return err
	}
	for _, b := range output.Buckets {
		fmt.Fprintf(out, "s3://%s/\n", *b.Name)
	}
	return nil
}

func iterateKeys(conn s3iface.S3API, urls []string, callback func(file File) error) error {
	found := false
	for _, url := range urls {
		fs := getFilesystem(conn, url)
		ch := fs.Files()
		for file := range ch {
			found = true
			err := callback(file)
			if err != nil {
				return err
			}
		}
		if fs.Error() != nil {
			return fs.Error()
		}
	}
	if !found {
		return ErrNotFound
	}
	return nil
}

func iterateKeysParallel(conn s3iface.S3API, urls []string, callback func(file File) error) error {
	// create pool for processing
	var err error
	wg := sync.WaitGroup{}
	q := make(chan File, 1000)
	for i := 0; i < parallel; i += 1 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for key := range q {
				e := callback(key)
				if e != nil {
					err = e
					return
				}
			}
		}()
	}

	e := iterateKeys(conn, urls, func(file File) error {
		q <- file
		return nil
	})
	if e != nil {
		return e
	}

	close(q)
	wg.Wait()
	return err
}

func listKeys(conn s3iface.S3API, urls []string) error {
	var count, totalSize int64
	err := iterateKeys(conn, urls, func(file File) error {
		if quiet {
			fmt.Fprintln(out, file)
		} else {
			fmt.Fprintf(out, "%s\t%db\n", file, file.Size())
		}
		count += 1
		totalSize += file.Size()
		return nil
	})
	if err != nil && err != ErrNotFound {
		return err
	}
	if !quiet {
		fmt.Fprintf(out, "\n%d files, %d bytes\n", count, totalSize)
	}
	return nil
}

func getKeys(conn s3iface.S3API, urls []string) error {
	for _, url := range urls {
		if !isS3Url(url) {
			return errors.New("s3:// url required")
		}
	}

	err := iterateKeysParallel(conn, urls, func(file File) error {
		reader, err := file.Reader()
		if err != nil {
			return err
		}
		defer reader.Close()

		// write files under relative path to the source path
		fpath := file.Relative()
		dirpath := path.Dir(fpath)
		if dirpath != "." {
			err = os.MkdirAll(dirpath, 0777)
			if err != nil {
				return err
			}
		}

		writer, err := os.Create(fpath)
		if err != nil {
			return err
		}
		nbytes, err := io.Copy(writer, reader)
		if err != nil {
			return err
		}
		if !quiet {
			fmt.Fprintf(out, "%s -> %s (%d bytes)\n", file, fpath, nbytes)
		}
		return nil
	})
	return err
}

func catKeys(conn s3iface.S3API, urls []string) error {
	return iterateKeysParallel(conn, urls, func(file File) error {
		reader, err := file.Reader()
		if err != nil {
			return err
		}
		defer reader.Close()

		if strings.HasSuffix(file.String(), ".gz") {
			reader, err = gzip.NewReader(reader)
			if err != nil {
				return err
			}
		}

		_, err = io.Copy(out, reader)
		if err != nil {
			return err
		}
		return nil
	})
}

func outputMatches(buf []byte, needle []byte, prefix string) {
	p := 0
	for {
		i := bytes.Index(buf[p:], needle)
		if i == -1 {
			break
		}
		i += p
		lineStart := bytes.LastIndexByte(buf[:i], byte('\n'))
		if lineStart == -1 {
			lineStart = 0
		} else {
			lineStart += 1
		}
		lineEnd := bytes.IndexByte(buf[i:], byte('\n'))
		if lineEnd == -1 {
			lineEnd = len(buf)
		} else {
			lineEnd += i
		}
		line := string(buf[lineStart:lineEnd])

		fmt.Fprintf(out, "%s%s\n", prefix, line)

		p = lineEnd + 1
		if p > len(buf)-len(needle) {
			break
		}
	}
}

func grepKeys(conn s3iface.S3API, find string, urls []string, noKeysPrefix bool, keysWithMatches bool) error {
	needle := []byte(find)

	return iterateKeysParallel(conn, urls, func(file File) error {
		reader, err := file.Reader()
		if err != nil {
			return err
		}
		defer reader.Close()

		if strings.HasSuffix(file.String(), ".gz") {
			reader, err = gzip.NewReader(reader)
			if err != nil {
				return err
			}
		}

		prefix := ""
		if !noKeysPrefix {
			prefix = file.String() + ":"
		}

		buf := make([]byte, 4096)
		offset := 0
		for n, err := reader.Read(buf[offset:]); n > 0 && err == nil; n, err = reader.Read(buf[offset:]) {
			if bytes.Contains(buf[:n+offset], needle) {
				if keysWithMatches {
					// only filename required, bail early
					fmt.Fprintln(out, file.String())
					break
				} else {
					outputMatches(buf[:n+offset], needle, prefix)
				}
			}
			// handle overlapping matches - copy last N-1 bytes to start of next
			offset = len(needle) - 1
			if offset > n {
				offset = n
			} else {
				copy(buf, buf[n-offset:])
			}
		}
{'fraction_non_alphanumeric': 0.15991034885987138, 'fraction_numerical': 0.011693626973299552, 'mean_word_length': 2.8039288361749444, 'pattern_counts': {'":': 2, '<': 9, '<?xml version=': 0, '>': 13, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5576973', 'n_tokens_mistral': 4491, 'n_tokens_neox': 3951, 'n_words': 1043}