takeEnd

CHANGES.md

@@ -15,6 +15,7 @@ To be released.  Unreleased versions are available on [nest.land].
  -  Added *src/range.ts* module.
  -  Added `range()` function.
  -  Added `Range` class.
+ -  Added `takeEnd()` function.
 
 [available on npm]: https://www.npmjs.com/package/aitertools
 

src/mod.ts

@@ -6,7 +6,7 @@ export { reduce } from "./fold.ts";
 export { count, cycle, repeat } from "./infinite.ts";
 export { map } from "./map.ts";
 export { Range, range } from "./range.ts";
-export { take, takeWhile } from "./take.ts";
+export { take, takeEnd, takeWhile } from "./take.ts";
 export { tee } from "./tee.ts";
 export { assertStreams, assertStreamStartsWith } from "./testing.ts";
 export { groupBy, unique } from "./unique.ts";

src/take.ts

@@ -133,3 +133,83 @@ export async function* takeWhile<T>(
     yield value;
   }
 }
+
+/**
+ * Takes a specified number of elements from the end of an async iterable.
+ *
+ * ``` typescript
+ * import { takeEnd } from "./take.ts";
+ * import { range } from "./range.ts";
+ *
+ * const iterable = takeEnd(range(10), 3);
+ * for await (const value of iterable) console.log(value);
+ * ```
+ *
+ * The above example will print the following 3 lines:
+ *
+ * ~~~
+ * 7
+ * 8
+ * 9
+ * ~~~
+ *
+ * If the iterable is shorter than the specified number, the whole elements are
+ * taken.
+ *
+ * ``` typescript
+ * import { takeEnd } from "./take.ts";
+ *
+ * async function* gen() { yield "foo"; yield "bar"; yield "baz"; }
+ * const iterable = takeEnd(gen(), 5);
+ * for await (const value of iterable) console.log(value);
+ * ```
+ *
+ * The above example will print only 3 elements, because `gen()` yields only 3
+ * elements:
+ *
+ * ~~~
+ * foo
+ * bar
+ * baz
+ * ~~~
+ *
+ * @template T The type of the elements in the `source` and the returned async
+ *             iterable.
+ * @param source The async iterable to take elements from.  It can be either
+ *               finite or infinite.
+ * @param count The number of elements to take.  It must be a finite integer.
+ * @returns An async iterable that yields the last `count` elements
+ *          from the `source` iterable.
+ * @throws `RangeError` if `count` is not a finite integer.
+ */
+export async function* takeEnd<T>(
+  source: Iterable<T> | AsyncIterable<T>,
+  count: number,
+): AsyncIterableIterator<T> {
+  if (!Number.isFinite(count)) throw new RangeError("count must be finite");
+  else if (!Number.isInteger(count)) {
+    throw new RangeError("count must be integer");
+  } else if (count < 1) return;
+
+  if (Array.isArray(source)) {
+    const length = source.length;
+    for (let i = Math.max(0, length - count); i < length; i++) {
+      yield source[i];
+    }
+    return;
+  }
+
+  const buffer: T[] = [];
+  let rewindPos = 0;
+  for await (const value of source) {
+    if (buffer.length >= count) {
+      buffer[rewindPos] = value;
+      if (rewindPos >= count - 1) rewindPos = 0;
+      else rewindPos++;
+    } else {
+      buffer.push(value);
+    }
+  }
+  for (let i = rewindPos; i < buffer.length; i++) yield buffer[i];
+  for (let i = 0; i < rewindPos; i++) yield buffer[i];
+}

tests/take_test.ts

@@ -1,8 +1,11 @@
-import { assert } from "https://deno.land/[email protected]/assert/mod.ts";
+import {
+  assert,
+  assertRejects,
+} from "https://deno.land/[email protected]/assert/mod.ts";
 import * as fc from "npm:[email protected]";
-import { fromIterable } from "../src/collections.ts";
+import { fromIterable, toArray } from "../src/collections.ts";
 import { count } from "../src/infinite.ts";
-import { take, takeWhile } from "../src/take.ts";
+import { take, takeEnd, takeWhile } from "../src/take.ts";
 import { assertStreams } from "../src/testing.ts";
 
 Deno.test("take()", async () => {
@@ -74,3 +77,56 @@ Deno.test("takeWhile() [fc]", async () => {
     );
   }
 });
+
+Deno.test("takeEnd() [fc]", async () => {
+  await fc.assert(
+    fc.asyncProperty(
+      fc.oneof(fc.array(fc.integer()), fc.array(fc.string())),
+      fc.integer(),
+      async (source: unknown[], count: number) => {
+        await assertStreams(
+          takeEnd(source, count),
+          count > 0 ? source.slice(-count) : [],
+        );
+        await assertStreams(
+          takeEnd(fromIterable(source), count),
+          count > 0 ? source.slice(-count) : [],
+        );
+      },
+    ),
+  );
+
+  await fc.assert(
+    fc.asyncProperty(
+      fc.oneof(fc.array(fc.integer()), fc.array(fc.string())),
+      fc.constantFrom(Infinity, -Infinity, NaN),
+      async (source: unknown[], count: number) => {
+        await assertRejects(
+          () => toArray(takeEnd(source, count)),
+          RangeError,
+        );
+        await assertRejects(
+          () => toArray(takeEnd(fromIterable(source), count)),
+          RangeError,
+        );
+      },
+    ),
+  );
+
+  await fc.assert(
+    fc.asyncProperty(
+      fc.oneof(fc.array(fc.integer()), fc.array(fc.string())),
+      fc.double().filter((n) => !Number.isInteger(n)),
+      async (source: unknown[], count: number) => {
+        await assertRejects(
+          () => toArray(takeEnd(source, count)),
+          RangeError,
+        );
+        await assertRejects(
+          () => toArray(takeEnd(fromIterable(source), count)),
+          RangeError,
+        );
+      },
+    ),
+  );
+});