geninq/src/functions/take.ts

A JavaScript version of the Linq library for Generators

import { Build } from "./utils";
import "./count";

declare global {
  interface Generator<T = unknown, TReturn = any, TNext = unknown> {
    take(count: number, at?: "last"): Generator<T, TReturn, TNext>;
  }

  interface AsyncGenerator<T = unknown, TReturn = any, TNext = unknown> {
    take(count: number, at?: "last"): AsyncGenerator<T, TReturn, TNext>;
  }
}

Build(
  "take",
  function* (count, at) {
    let items: Iterable<unknown> = this;
    if (at === "last") {
      const array = this.array();
      count = array.length - count;
      items = array;
    }

    let index = 0;
    for (const item of items) {
      if (!at && index < count) {
        yield item;
      } else if (at === "last" && index >= count) {
        yield item;
      }

      index++;
    }
  },
  async function* (count, at) {
    let items: AsyncIterable<unknown> | Iterable<unknown> = this;
    if (at === "last") {
      const array = await this.array();
      count = array.length - count;
      items = array;
    }

    let index = 0;
    for await (const item of items) {
      if (!at && index < count) {
        yield item;
      } else if (at === "last" && index >= count) {
        yield item;
      }

      index++;
    }
  }
);