feat: add types to keyboard controls

.storybook/stories/KeyboardControls.stories.tsx

@@ -0,0 +1,72 @@
+import { useFrame } from '@react-three/fiber'
+import * as React from 'react'
+import { useMemo, useRef } from 'react'
+import { MathUtils, Mesh, Vector3 } from 'three'
+import { Cone, KeyboardControls, KeyboardControlsEntry, useKeyboardControls } from '../../src'
+import { Setup } from '../Setup'
+
+export default {
+  title: 'Controls/KeyboardControls',
+  decorators: [
+    (storyFn) => (
+      <Setup cameraPosition={new Vector3(0, 10, 0)} lights={true}>
+        {storyFn()}
+      </Setup>
+    ),
+  ],
+}
+
+enum Controls {
+  forward = 'forward',
+  left = 'left',
+  right = 'right',
+  back = 'back',
+}
+
+export const KeyboardControlsSt = () => {
+  const map = useMemo<KeyboardControlsEntry[]>(
+    () => [
+      { name: Controls.forward, keys: ['ArrowUp', 'w', 'W'] },
+      { name: Controls.back, keys: ['ArrowDown', 's', 'S'] },
+      { name: Controls.left, keys: ['ArrowLeft', 'a', 'A'] },
+      { name: Controls.right, keys: ['ArrowRight', 'd', 'D'] },
+    ],
+    []
+  )
+
+  return (
+    <KeyboardControls map={map}>
+      <Player />
+    </KeyboardControls>
+  )
+}
+
+const _velocity = new Vector3()
+const speed = 10
+
+const Player = () => {
+  const ref = useRef<Mesh>(null)
+  const [, get] = useKeyboardControls<Controls>()
+
+  useFrame((_s, dl) => {
+    if (!ref.current) return
+    const state = get()
+    if (state.left && !state.right) _velocity.x = -1
+    if (state.right && !state.left) _velocity.x = 1
+    if (!state.left && !state.right) _velocity.x = 0
+
+    if (state.forward && !state.back) _velocity.z = -1
+    if (state.back && !state.forward) _velocity.z = 1
+    if (!state.forward && !state.back) _velocity.z = 0
+
+    ref.current.position.addScaledVector(_velocity, speed * dl)
+
+    ref.current.rotateY(4 * dl * _velocity.x)
+  })
+
+  return (
+    <Cone ref={ref} args={[1, 3, 4]} rotation={[-90 * MathUtils.DEG2RAD, 0, 0]}>
+      <meshLambertMaterial color="green" />
+    </Cone>
+  )
+}

src/web/KeyboardControls.tsx

@@ -2,9 +2,9 @@ import * as React from 'react'
 import create, { GetState, StateSelector, Subscribe, UseBoundStore } from 'zustand'
 import { subscribeWithSelector } from 'zustand/middleware'
 
-type KeyboardControlsState = { [key: string]: boolean }
+type KeyboardControlsState<T extends string = string> = { [K in T]: boolean }
 
-type KeyboardControlsEntry = {
+export type KeyboardControlsEntry = {
   /** Name of the action */
   name: string
   /** The keys that define it, you can use either event.key, or event.code */
@@ -24,10 +24,10 @@ type KeyboardControlsProps = {
   domElement?: HTMLElement
 }
 
-type KeyboardControlsApi = [
-  Subscribe<KeyboardControlsState>,
-  GetState<KeyboardControlsState>,
-  UseBoundStore<KeyboardControlsState>
+type KeyboardControlsApi<T extends string = string> = [
+  Subscribe<KeyboardControlsState<T>>,
+  GetState<KeyboardControlsState<T>>,
+  UseBoundStore<KeyboardControlsState<T>>
 ]
 
 const context = /*@__PURE__*/ React.createContext<KeyboardControlsApi>(null!)
@@ -90,8 +90,17 @@ export function KeyboardControls({ map, children, onChange, domElement }: Keyboa
   return <context.Provider value={api} children={children} />
 }
 
-export function useKeyboardControls(sel?: StateSelector<KeyboardControlsState, any>) {
-  const [sub, get, store] = React.useContext(context)
+export function useKeyboardControls<T extends string = string, U = any>(): [
+  Subscribe<KeyboardControlsState<T>>,
+  GetState<KeyboardControlsState<T>>
+]
+export function useKeyboardControls<T extends string = string, U = any>(
+  sel: StateSelector<KeyboardControlsState<T>, U>
+): U
+export function useKeyboardControls<T extends string = string, U = any>(
+  sel?: StateSelector<KeyboardControlsState<T>, U>
+): U | [Subscribe<KeyboardControlsState<T>>, GetState<KeyboardControlsState<T>>] {
+  const [sub, get, store] = React.useContext<KeyboardControlsApi<T>>(context)
   if (sel) return store(sel)
   else return [sub, get]
 }