From f586d06759aff8099f5602f56803cbce1f25cbf3 Mon Sep 17 00:00:00 2001
From: Anton Golub <mailbox@antongolub.ru>
Date: Fri, 18 Mar 2022 11:16:02 +0300
Subject: [PATCH] feat: introduce experimental `withTimeout` helper

closes #349
---
 README.md             | 10 ++++++++++
 src/experimental.d.ts |  9 ++++-----
 src/experimental.mjs  | 10 ++++++++++
 src/index.d.ts        |  4 +++-
 test.mjs              | 15 ++++++++++++++-
 5 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 02b9eb1084..ded79802d7 100644
--- a/README.md
+++ b/README.md
@@ -408,6 +408,16 @@ await $`long-running command`
 stop()
 ```
 
+#### `withTimeout()`
+
+Runs and sets a timeout for a cmd.
+
+```js
+import {withTimeout} from 'zx/experimental'
+
+await withTimeout(100, 'SIGTERM')`sleep 9999`
+```
+
 ### FAQ
 
 #### Passing env variables
diff --git a/src/experimental.d.ts b/src/experimental.d.ts
index 5e4b1a51a2..9ec1eacc0b 100644
--- a/src/experimental.d.ts
+++ b/src/experimental.d.ts
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import {ProcessOutput} from './index'
+import {ZxTemplate} from './index'
 
 interface Echo {
   (pieces: TemplateStringsArray, ...args: any[]): void
@@ -20,10 +20,9 @@ interface Echo {
 }
 export const echo: Echo
 
-interface Retry {
-  (pieces: TemplateStringsArray, ...args: any[]): Promise<ProcessOutput>
-}
-export const retry: (count?: number, delay?: number) => Retry
+export const retry: (count?: number, delay?: number) => ZxTemplate
+
+export const withTimeout: (delay?: number, signal?: string | number) => ZxTemplate
 
 type StopSpinner = () => void
 export function startSpinner(title: string): StopSpinner
diff --git a/src/experimental.mjs b/src/experimental.mjs
index ffd15aa045..0bab6ad0a9 100644
--- a/src/experimental.mjs
+++ b/src/experimental.mjs
@@ -25,6 +25,16 @@ export const retry = (count = 5, delay = 0) => async (cmd, ...args) => {
   }
 }
 
+// Runs and sets a timeout for a cmd
+export const withTimeout = (timeout, signal) => async (cmd, ...args) => {
+  let p = $(cmd, ...args)
+  if (!timeout) return p
+
+  let timer = setTimeout(() => p.kill(signal), timeout)
+
+  return p.finally(() => clearTimeout(timer))
+}
+
 // A console.log() alternative which can take ProcessOutput.
 export function echo(pieces, ...args) {
   if (Array.isArray(pieces) && pieces.every(isString) && pieces.length - 1 === args.length) {
diff --git a/src/index.d.ts b/src/index.d.ts
index 0d99c20ad6..42ff8cf6cb 100644
--- a/src/index.d.ts
+++ b/src/index.d.ts
@@ -23,9 +23,11 @@ import * as _yaml from 'yaml'
 import _fetch from 'node-fetch'
 import {ParsedArgs} from 'minimist'
 
-interface $ {
+export interface ZxTemplate {
   (pieces: TemplateStringsArray, ...args: any[]): ProcessPromise<ProcessOutput>
+}
 
+interface $ extends ZxTemplate {
   verbose: boolean
   shell: string
   prefix: string
diff --git a/test.mjs b/test.mjs
index 3e64eba5a5..fd42a86f5a 100755
--- a/test.mjs
+++ b/test.mjs
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 import {strict as assert} from 'assert'
-import {retry} from './src/experimental.mjs'
+import {retry, withTimeout} from './src/experimental.mjs'
 
 let всегоТестов = 0
 
@@ -257,6 +257,19 @@ if (test('Retry works')) {
   assert(Date.now() >= now + 50 * (5 - 1))
 }
 
+if (test('withTimeout works')) {
+  let exitCode = 0
+  let signal
+  try {
+    await withTimeout(100, 'SIGKILL')`sleep 9999`
+  } catch (p) {
+    exitCode = p.exitCode
+    signal = p.signal
+  }
+  assert.equal(exitCode, null)
+  assert.equal(signal, 'SIGKILL')
+}
+
 let version
 if (test('require() is working in ESM')) {
   let data = require('./package.json')