From 22fbca58d34f2f2d6c12e51a34244af4227b8885 Mon Sep 17 00:00:00 2001
From: Steve Ruiz <steveruizok@gmail.com>
Date: Mon, 13 Sep 2021 16:59:37 +0100
Subject: [PATCH] Allows for additional properties

---
 .../core/src/components/renderer/renderer.tsx |  2 +-
 packages/core/src/shapes/createShape.spec.tsx |  7 ++++++-
 packages/core/src/shapes/createShape.tsx      | 21 ++++++++++++-------
 packages/tldraw/src/shape/shape-utils.tsx     |  4 ++--
 .../session/sessions/brush/brush.session.ts   |  1 -
 packages/tldraw/src/state/tldr.ts             |  4 ++--
 packages/tldraw/src/types.ts                  |  3 +--
 7 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/packages/core/src/components/renderer/renderer.tsx b/packages/core/src/components/renderer/renderer.tsx
index 58c08270d..f29d8d6d1 100644
--- a/packages/core/src/components/renderer/renderer.tsx
+++ b/packages/core/src/components/renderer/renderer.tsx
@@ -14,7 +14,7 @@ import { Inputs } from '../../inputs'
 import { useTLTheme, TLContext, TLContextType } from '../../hooks'
 import type { TLShapeUtil } from '+index'
 
-export interface RendererProps<T extends TLShape, E extends Element, M = any>
+export interface RendererProps<T extends TLShape, E extends Element = any, M = any>
   extends Partial<TLCallbacks<T>> {
   /**
    * An object containing instances of your shape classes.
diff --git a/packages/core/src/shapes/createShape.spec.tsx b/packages/core/src/shapes/createShape.spec.tsx
index 4217f3016..133bf7fcf 100644
--- a/packages/core/src/shapes/createShape.spec.tsx
+++ b/packages/core/src/shapes/createShape.spec.tsx
@@ -12,8 +12,9 @@ export interface BoxShape extends TLShape {
   size: number[]
 }
 
-export const Box = new ShapeUtil<BoxShape, SVGSVGElement, null>(() => {
+export const Box = new ShapeUtil<BoxShape, SVGSVGElement, null, { age: number }>(() => {
   return {
+    age: 100,
     type: 'box',
     defaultProps: {
       id: 'example1',
@@ -90,6 +91,10 @@ describe('shape utils', () => {
     expect(Box).toBeTruthy()
   })
 
+  it('creates a shape utils with extended properties', () => {
+    expect(Box.age).toBe(100)
+  })
+
   it('creates a shape', () => {
     expect(Box.create({ id: 'box1' })).toStrictEqual({
       ...boxShape,
diff --git a/packages/core/src/shapes/createShape.tsx b/packages/core/src/shapes/createShape.tsx
index d5bcf5c2e..f9d274f2d 100644
--- a/packages/core/src/shapes/createShape.tsx
+++ b/packages/core/src/shapes/createShape.tsx
@@ -5,13 +5,14 @@ import type { TLShape, TLShapeUtil } from '+types'
 import Utils from '+utils'
 import { intersectPolylineBounds, intersectRayBounds } from '@tldraw/intersect'
 
-export const ShapeUtil = function <T extends TLShape, E extends Element, M = any>(
-  this: TLShapeUtil<T, E, M>,
+export const ShapeUtil = function <T extends TLShape, E extends Element, M = any, K = unknown>(
+  this: TLShapeUtil<T, E, M> & K,
   fn: (
-    this: TLShapeUtil<T, E, M>
+    this: TLShapeUtil<T, E, M> & K
   ) => Partial<TLShapeUtil<T, E, M>> &
-    Pick<TLShapeUtil<T, E, M>, 'type' | 'defaultProps' | 'Component' | 'Indicator' | 'getBounds'>
-) {
+    Pick<TLShapeUtil<T, E, M>, 'type' | 'defaultProps' | 'Component' | 'Indicator' | 'getBounds'> &
+    K
+): TLShapeUtil<T, E, M> & ReturnType<typeof fn> {
   const defaults: Partial<TLShapeUtil<T, E, M>> = {
     refMap: new Map(),
 
@@ -211,10 +212,14 @@ export const ShapeUtil = function <T extends TLShape, E extends Element, M = any
 
   return this
 } as unknown as {
-  new <T extends TLShape, E extends Element, M = any>(
+  new <T extends TLShape, E extends Element, M = any, K = unknown>(
     fn: (
       this: TLShapeUtil<T, E, M>
     ) => Partial<TLShapeUtil<T, E, M>> &
-      Pick<TLShapeUtil<T, E, M>, 'type' | 'defaultProps' | 'Component' | 'Indicator' | 'getBounds'>
-  ): TLShapeUtil<T, E, M>
+      Pick<
+        TLShapeUtil<T, E, M>,
+        'type' | 'defaultProps' | 'Component' | 'Indicator' | 'getBounds'
+      > &
+      K
+  ): TLShapeUtil<T, E, M> & ReturnType<typeof fn>
 }
diff --git a/packages/tldraw/src/shape/shape-utils.tsx b/packages/tldraw/src/shape/shape-utils.tsx
index 4e05016dc..5f0faab6e 100644
--- a/packages/tldraw/src/shape/shape-utils.tsx
+++ b/packages/tldraw/src/shape/shape-utils.tsx
@@ -13,6 +13,6 @@ export const tldrawShapeUtils: Record<TLDrawShapeType, any> = {
   [TLDrawShapeType.PostIt]: PostIt,
 }
 
-export function getShapeUtils<T extends TLDrawShape>(type: TLDrawShapeType) {
-  return tldrawShapeUtils[type] as TLDrawShapeUtil<T, any>
+export function getShapeUtils<T extends TLDrawShape>(type: T['type']) {
+  return tldrawShapeUtils[type] as TLDrawShapeUtil<T>
 }
diff --git a/packages/tldraw/src/state/session/sessions/brush/brush.session.ts b/packages/tldraw/src/state/session/sessions/brush/brush.session.ts
index 58a4a3ec4..01b151594 100644
--- a/packages/tldraw/src/state/session/sessions/brush/brush.session.ts
+++ b/packages/tldraw/src/state/session/sessions/brush/brush.session.ts
@@ -1,7 +1,6 @@
 import { brushUpdater, Utils } from '@tldraw/core'
 import { Vec } from '@tldraw/vec'
 import { Data, Session, TLDrawPatch, TLDrawStatus } from '~types'
-import { getShapeUtils } from '~shape'
 import { TLDR } from '~state/tldr'
 
 export class BrushSession implements Session {
diff --git a/packages/tldraw/src/state/tldr.ts b/packages/tldraw/src/state/tldr.ts
index eebe239b7..b6de71bbe 100644
--- a/packages/tldraw/src/state/tldr.ts
+++ b/packages/tldraw/src/state/tldr.ts
@@ -15,9 +15,9 @@ import { Vec } from '@tldraw/vec'
 
 export class TLDR {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  static getShapeUtils<T extends TLDrawShape>(type: T['type']): TLDrawShapeUtil<T, any>
+  static getShapeUtils<T extends TLDrawShape>(type: T['type']): TLDrawShapeUtil<T>
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  static getShapeUtils<T extends TLDrawShape>(shape: T): TLDrawShapeUtil<T, any>
+  static getShapeUtils<T extends TLDrawShape>(shape: T): TLDrawShapeUtil<T>
   static getShapeUtils<T extends TLDrawShape>(shape: T | T['type']) {
     return getShapeUtils<T>(typeof shape === 'string' ? shape : shape.type)
   }
diff --git a/packages/tldraw/src/types.ts b/packages/tldraw/src/types.ts
index 946357a82..5d8caedd8 100644
--- a/packages/tldraw/src/types.ts
+++ b/packages/tldraw/src/types.ts
@@ -212,8 +212,7 @@ export type TLDrawShape =
   | GroupShape
   | PostItShape
 
-export interface TLDrawShapeUtil<T extends TLDrawShape, E extends Element>
-  extends TLShapeUtil<T, E, TLDrawMeta> {
+export type TLDrawShapeUtil<T extends TLDrawShape> = TLShapeUtil<T, any, TLDrawMeta> & {
   toolType: TLDrawToolType
 }