Skip to content

all: simplify iterator API #13

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
changkun merged 1 commit into from
Dec 21, 2025
Merged

all: simplify iterator API #13

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 8 additions & 6 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,12 +194,14 @@ result, err := c.GetParseResult(
tensorlake.WithSSE(true),
tensorlake.WithOnUpdate(func(name tensorlake.ParseEventName, r *tensorlake.ParseResult) {
switch eventName {
case tensorlake.sseEventParseQueued:
case tensorlake.SSEEventParseQueued:
fmt.Println("Job queued")
case tensorlake.sseEventParseUpdate:
case tensorlake.SSEEventParseUpdate:
fmt.Printf("Progress: %d/%d pages\n", r.ParsedPagesCount, r.TotalPages)
case tensorlake.sseEventParseDone:
case tensorlake.SSEEventParseDone:
fmt.Println("Complete!")
case tensorlake.SSEEventParseFailed:
fmt.Printf("Failed: %s\n", r.Error)
}
}),
)
Expand All @@ -211,23 +213,23 @@ Easily iterate through paginated results:

```go
// Iterate all files
for file, err := range c.IterFiles(ctx, 50, tensorlake.PaginationDirectionNext) {
for file, err := range c.IterFiles(ctx, 50) {
if err != nil {
panic(err)
}
fmt.Printf("File: %s\n", file.FileName)
}

// Iterate all parse jobs
for job, err := range c.IterParseJobs(ctx, 50, tensorlake.PaginationDirectionNext) {
for job, err := range c.IterParseJobs(ctx, 50) {
if err != nil {
panic(err)
}
fmt.Printf("Job %s: Status: %s\n", job.ParseId, job.Status)
}

// Iterate all datasets
for dataset, err := range c.IterDatasets(ctx, 50, tensorlake.PaginationDirectionNext) {
for dataset, err := range c.IterDatasets(ctx, 50) {
if err != nil {
panic(err)
}
Expand Down
13 changes: 7 additions & 6 deletions dataset_list.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,14 @@ import (
)

// IterDatasets iterates over all datasets in the organization.
func (c *Client) IterDatasets(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[Dataset, error] {
func (c *Client) IterDatasets(ctx context.Context, batchSize int) iter.Seq2[Dataset, error] {
return func(yield func(Dataset, error) bool) {
cursor := ""
for {
listResp, err := c.ListDatasets(ctx, &ListDatasetsRequest{
Cursor: cursor,
Limit: limit,
Direction: direction,
Limit: batchSize,
Direction: PaginationDirectionNext,
})
if err != nil {
yield(Dataset{}, err)
Expand All @@ -52,14 +52,15 @@ func (c *Client) IterDatasets(ctx context.Context, limit int, direction Paginati
}

// IterDatasetData iterates over all dataset data in the organization.
func (c *Client) IterDatasetData(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[ParseResult, error] {
func (c *Client) IterDatasetData(ctx context.Context, datasetId string, batchSize int) iter.Seq2[ParseResult, error] {
return func(yield func(ParseResult, error) bool) {
cursor := ""
for {
listResp, err := c.ListDatasetData(ctx, &ListDatasetDataRequest{
DatasetId: datasetId,
Cursor: cursor,
Limit: limit,
Direction: direction,
Limit: batchSize,
Direction: PaginationDirectionNext,
})
if err != nil {
yield(ParseResult{}, err)
Expand Down
4 changes: 2 additions & 2 deletions dataset_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ func TestDataset(t *testing.T) {

// List datasets.
datasets := []string{}
for d, err := range c.IterDatasets(t.Context(), 1, PaginationDirectionNext) {
for d, err := range c.IterDatasets(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list datasets: %v", err)
}
Expand Down Expand Up @@ -113,7 +113,7 @@ func TestDataset(t *testing.T) {

// Check if the dataset is deleted.
datasets = []string{}
for d, err := range c.IterDatasets(t.Context(), 1, PaginationDirectionNext) {
for d, err := range c.IterDatasets(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list datasets: %v", err)
}
Expand Down
4 changes: 2 additions & 2 deletions docs/dataset-apis.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -285,13 +285,13 @@ if response.HasMore {
For convenience, use the `IterDatasets` method:

```go
func (c *Client) IterDatasets(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[Dataset, error]
func (c *Client) IterDatasets(ctx context.Context, batchSize int) iter.Seq2[Dataset, error]
```

**Example:**

```go
for dataset, err := range client.IterDatasets(context.Background(), 50, tensorlake.PaginationDirectionNext) {
for dataset, err := range client.IterDatasets(context.Background(), 50) {
if err != nil {
log.Fatal(err)
}
Expand Down
4 changes: 2 additions & 2 deletions docs/file-apis.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,13 +175,13 @@ if response.HasMore {
For convenience, use the `IterFiles` method to iterate through all files:

```go
func (c *Client) IterFiles(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[FileInfo, error]
func (c *Client) IterFiles(ctx context.Context, batchSize int) iter.Seq2[FileInfo, error]
```

**Example:**

```go
for file, err := range client.IterFiles(context.Background(), 50, tensorlake.PaginationDirectionNext) {
for file, err := range client.IterFiles(context.Background(), 50) {
if err != nil {
log.Fatal(err)
}
Expand Down
4 changes: 2 additions & 2 deletions docs/parse-apis.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -453,13 +453,13 @@ for _, job := range response.Items {
### Iterate All Parse Jobs

```go
func (c *Client) IterParseJobs(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[ParseResult, error]
func (c *Client) IterParseJobs(ctx context.Context, batchSize int) iter.Seq2[ParseResult, error]
```

**Example:**

```go
for job, err := range client.IterParseJobs(context.Background(), 50, tensorlake.PaginationDirectionNext) {
for job, err := range client.IterParseJobs(context.Background(), 50) {
if err != nil {
log.Fatal(err)
}
Expand Down
6 changes: 3 additions & 3 deletions file_list.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,14 @@ import (
)

// IterFiles iterates over all files in the project.
func (c *Client) IterFiles(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[FileInfo, error] {
func (c *Client) IterFiles(ctx context.Context, batchSize int) iter.Seq2[FileInfo, error] {
return func(yield func(FileInfo, error) bool) {
cursor := ""
for {
listResp, err := c.ListFiles(ctx, &ListFilesRequest{
Cursor: cursor,
Limit: limit,
Direction: direction,
Limit: batchSize,
Direction: PaginationDirectionNext,
})
if err != nil {
yield(FileInfo{}, err)
Expand Down
4 changes: 2 additions & 2 deletions file_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ func TestFileManagement(t *testing.T) {

// List the files. Iterate through all the pages.
files := []string{}
for f, err := range c.IterFiles(t.Context(), 1, PaginationDirectionNext) {
for f, err := range c.IterFiles(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list files: %v", err)
}
Expand Down Expand Up @@ -124,7 +124,7 @@ func TestFileManagement(t *testing.T) {

// Validate file is deleted.
files = []string{}
for f, err := range c.IterFiles(t.Context(), 1, PaginationDirectionNext) {
for f, err := range c.IterFiles(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list files: %v", err)
}
Expand Down
2 changes: 1 addition & 1 deletion parse_classify_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ func TestClassifyDocument(t *testing.T) {

// Validate classify results.
jobs := []string{}
for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) {
for j, err := range c.IterParseJobs(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list parse jobs: %v", err)
}
Expand Down
2 changes: 1 addition & 1 deletion parse_extract_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ func TestExtractDocument(t *testing.T) {

// Validate parse results.
jobs := []string{}
for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) {
for j, err := range c.IterParseJobs(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list parse jobs: %v", err)
}
Expand Down
14 changes: 7 additions & 7 deletions parse_get.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,10 +103,10 @@ type ParseEventName string
// The possible SSE events.
// See also: https://github.com/tensorlakeai/tensorlake/blob/main/src/tensorlake/documentai/_parse.py#L499
const (
sseEventParseQueued ParseEventName = "parse_queued"
sseEventParseUpdate ParseEventName = "parse_update"
sseEventParseDone ParseEventName = "parse_done"
sseEventParseFailed ParseEventName = "parse_failed"
SSEEventParseQueued ParseEventName = "parse_queued"
SSEEventParseUpdate ParseEventName = "parse_update"
SSEEventParseDone ParseEventName = "parse_done"
SSEEventParseFailed ParseEventName = "parse_failed"
)

func (c *Client) handleSSEResponse(req *http.Request, onUpdate ParseResultUpdateFunc) (*ParseResult, error) {
Expand Down Expand Up @@ -151,19 +151,19 @@ func (c *Client) handleSSEResponse(req *http.Request, onUpdate ParseResultUpdate
}

switch ev.Name() {
case string(sseEventParseQueued), string(sseEventParseUpdate):
case string(SSEEventParseQueued), string(SSEEventParseUpdate):
if onUpdate != nil {
onUpdate(ParseEventName(ev.Name()), &result)
}
continue

case string(sseEventParseDone):
case string(SSEEventParseDone):
if onUpdate != nil {
onUpdate(ParseEventName(ev.Name()), &result)
}
return &result, nil

case string(sseEventParseFailed):
case string(SSEEventParseFailed):
if onUpdate != nil {
onUpdate(ParseEventName(ev.Name()), &result)
}
Expand Down
2 changes: 1 addition & 1 deletion parse_get_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ func TestGetParseResultSSE(t *testing.T) {

// Validate parse results.
jobs := []string{}
for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) {
for j, err := range c.IterParseJobs(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list parse jobs: %v", err)
}
Expand Down
6 changes: 3 additions & 3 deletions parse_list.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,14 @@ import (
)

// IterParseJobs iterates over all parse jobs in the project.
func (c *Client) IterParseJobs(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[ParseResult, error] {
func (c *Client) IterParseJobs(ctx context.Context, batchSize int) iter.Seq2[ParseResult, error] {
return func(yield func(ParseResult, error) bool) {
cursor := ""
for {
listResp, err := c.ListParseJobs(ctx, &ListParseJobsRequest{
Cursor: cursor,
Limit: limit,
Direction: direction,
Limit: batchSize,
Direction: PaginationDirectionNext,
})
if err != nil {
yield(ParseResult{}, err)
Expand Down
4 changes: 2 additions & 2 deletions parse_parse_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ func TestParseDocumentRemote(t *testing.T) {

// Validate parse results.
jobs := []string{}
for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) {
for j, err := range c.IterParseJobs(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list parse jobs: %v", err)
}
Expand Down Expand Up @@ -212,7 +212,7 @@ func TestParseDocumentStructuredExtraction(t *testing.T) {

// Validate parse results.
jobs := []string{}
for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) {
for j, err := range c.IterParseJobs(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list parse jobs: %v", err)
}
Expand Down
4 changes: 2 additions & 2 deletions util_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ func testCleanupFileAndParseJob(t *testing.T, c *Client, fileId string, parseId

// Check if file is deleted.
files := []string{}
for f, err := range c.IterFiles(t.Context(), 1, PaginationDirectionNext) {
for f, err := range c.IterFiles(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list files: %v", err)
}
Expand All @@ -48,7 +48,7 @@ func testCleanupFileAndParseJob(t *testing.T, c *Client, fileId string, parseId

// Check if parse job is deleted.
jobs := []string{}
for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) {
for j, err := range c.IterParseJobs(t.Context(), 1) {
if err != nil {
t.Fatalf("failed to list parse jobs: %v", err)
}
Expand Down