From f7fcc75196243ba990ec0c0f5ad2f923f4cb0d65 Mon Sep 17 00:00:00 2001
From: Neil Kistner <neil.kistner@gmail.com>
Date: Mon, 7 Dec 2020 20:16:38 -0600
Subject: [PATCH] fix(axes): improve package types

---
 packages/axes/src/canvas.ts               | 21 ++++++++++------
 packages/axes/src/components/Axes.tsx     |  4 ++-
 packages/axes/src/components/Axis.tsx     |  8 +++---
 packages/axes/src/components/AxisTick.tsx |  2 +-
 packages/axes/src/compute.ts              | 20 +++++++++------
 packages/axes/src/types.ts                | 30 +++++++++++------------
 6 files changed, 48 insertions(+), 37 deletions(-)

diff --git a/packages/axes/src/canvas.ts b/packages/axes/src/canvas.ts
index c626b1e9a0..cf7435da7a 100644
--- a/packages/axes/src/canvas.ts
+++ b/packages/axes/src/canvas.ts
@@ -2,7 +2,7 @@ import { degreesToRadians, CompleteTheme } from '@nivo/core'
 import { computeCartesianTicks, getFormatter, computeGridLines } from './compute'
 import { TicksSpec, AnyScale, AxisLegendPosition, CanvasAxisProp, ValueFormatter } from './types'
 
-export const renderAxisToCanvas = <Value>(
+export const renderAxisToCanvas = <Value extends string | number | Date>(
     ctx: CanvasRenderingContext2D,
     {
         axis,
@@ -34,7 +34,7 @@ export const renderAxisToCanvas = <Value>(
         tickSize?: number
         tickPadding?: number
         tickRotation?: number
-        format?: ValueFormatter
+        format?: string | ValueFormatter<Value>
         legend?: string
         legendPosition?: AxisLegendPosition
         legendOffset?: number
@@ -88,7 +88,7 @@ export const renderAxisToCanvas = <Value>(
             ctx.stroke()
         }
 
-        const value = format !== undefined ? format(String(tick.value)) : (tick.value as string)
+        const value = typeof format === 'function' ? format(tick.value) : (tick.value as string)
 
         ctx.save()
         ctx.translate(tick.x + tick.textX, tick.y + tick.textY)
@@ -98,7 +98,7 @@ export const renderAxisToCanvas = <Value>(
             ctx.fillStyle = theme.axis.ticks.text.fill
         }
 
-        ctx.fillText(value, 0, 0)
+        ctx.fillText(String(value), 0, 0)
         ctx.restore()
     })
 
@@ -153,7 +153,10 @@ export const renderAxisToCanvas = <Value>(
 
 const positions = ['top', 'right', 'bottom', 'left'] as const
 
-export const renderAxesToCanvas = <X, Y>(
+export const renderAxesToCanvas = <
+    X extends string | number | Date,
+    Y extends string | number | Date
+>(
     ctx: CanvasRenderingContext2D,
     {
         xScale,
@@ -182,7 +185,9 @@ export const renderAxesToCanvas = <X, Y>(
     const axes = { top, right, bottom, left }
 
     positions.forEach(position => {
-        const axis = axes[position]
+        const axis = axes[position] as typeof position extends 'bottom' | 'top'
+            ? CanvasAxisProp<X> | undefined
+            : CanvasAxisProp<Y> | undefined
 
         if (!axis) return null
 
@@ -191,7 +196,7 @@ export const renderAxesToCanvas = <X, Y>(
         const scale = isXAxis ? xScale : yScale
         const format = getFormatter(axis.format, scale)
 
-        renderAxisToCanvas<unknown>(ctx, {
+        renderAxisToCanvas(ctx, {
             ...axis,
             axis: isXAxis ? 'x' : 'y',
             x: position === 'right' ? width : 0,
@@ -205,7 +210,7 @@ export const renderAxesToCanvas = <X, Y>(
     })
 }
 
-export const renderGridLinesToCanvas = <Value>(
+export const renderGridLinesToCanvas = <Value extends string | number | Date>(
     ctx: CanvasRenderingContext2D,
     {
         width,
diff --git a/packages/axes/src/components/Axes.tsx b/packages/axes/src/components/Axes.tsx
index 5f499b2f1d..d97a6bd05f 100644
--- a/packages/axes/src/components/Axes.tsx
+++ b/packages/axes/src/components/Axes.tsx
@@ -28,7 +28,9 @@ export const Axes = <X extends number | string | Date, Y extends number | string
     return (
         <>
             {positions.map(position => {
-                const axis = axes[position]
+                const axis = axes[position] as typeof position extends 'bottom' | 'top'
+                    ? AxisProp<X> | undefined
+                    : AxisProp<Y> | undefined
 
                 if (!axis) return null
 
diff --git a/packages/axes/src/components/Axis.tsx b/packages/axes/src/components/Axis.tsx
index fc17972c39..664ae2bf06 100644
--- a/packages/axes/src/components/Axis.tsx
+++ b/packages/axes/src/components/Axis.tsx
@@ -3,9 +3,9 @@ import { useSpring, useTransition, animated } from 'react-spring'
 import { useTheme, useMotionConfig } from '@nivo/core'
 import { computeCartesianTicks, getFormatter } from '../compute'
 import { AxisTick } from './AxisTick'
-import { AxisProps, AxisTickProps } from '../types'
+import { AxisProps } from '../types'
 
-export const Axis = <Scale,>({
+export const Axis = <Value extends string | number | Date>({
     axis,
     scale,
     x = 0,
@@ -23,14 +23,14 @@ export const Axis = <Scale,>({
     legendOffset = 0,
     onClick,
     ariaHidden,
-}: AxisProps) => {
+}: AxisProps<Value>) => {
     const theme = useTheme()
 
     const formatValue = useMemo(() => getFormatter(format, scale), [format, scale])
 
     const { ticks, textAlign, textBaseline } = computeCartesianTicks({
         axis,
-        scale: scale as any,
+        scale,
         ticksPosition,
         tickValues,
         tickSize,
diff --git a/packages/axes/src/components/AxisTick.tsx b/packages/axes/src/components/AxisTick.tsx
index 6407b93a15..adcc35bbff 100644
--- a/packages/axes/src/components/AxisTick.tsx
+++ b/packages/axes/src/components/AxisTick.tsx
@@ -16,7 +16,7 @@ export const AxisTick = <Value extends string | number | Date>({
     const theme = useTheme()
 
     let value = _value
-    if (format !== undefined) {
+    if (typeof format === 'function') {
         value = format(value)
     }
 
diff --git a/packages/axes/src/compute.ts b/packages/axes/src/compute.ts
index 6f601858b0..149c5f4c5e 100644
--- a/packages/axes/src/compute.ts
+++ b/packages/axes/src/compute.ts
@@ -76,7 +76,10 @@ const isInteger = (value: unknown): value is number =>
 
 const isArray = <T>(value: unknown): value is T[] => Array.isArray(value)
 
-export const getScaleTicks = <Value>(scale: AnyScale, spec?: TicksSpec<Value>) => {
+export const getScaleTicks = <Value extends string | number | Date>(
+    scale: AnyScale,
+    spec?: TicksSpec<Value>
+) => {
     // specific values
     if (Array.isArray(spec)) {
         return spec
@@ -122,7 +125,7 @@ export const getScaleTicks = <Value>(scale: AnyScale, spec?: TicksSpec<Value>) =
     return scale.domain()
 }
 
-export const computeCartesianTicks = <Value>({
+export const computeCartesianTicks = <Value extends string | number | Date>({
     axis,
     scale,
     ticksPosition,
@@ -209,21 +212,22 @@ export const computeCartesianTicks = <Value>({
     }
 }
 
-export const getFormatter = (
-    format: string | ValueFormatter,
+export const getFormatter = <Value extends string | number | Date>(
+    format: string | ValueFormatter<Value> | undefined,
     scale: AnyScale
-): ValueFormatter | undefined => {
+): ValueFormatter<Value> | undefined => {
     if (typeof format === 'undefined' || typeof format === 'function') return format
 
     if (scale.type === 'time') {
         const formatter = timeFormat(format)
-        return (d: number | string) => formatter(new Date(d))
+
+        return (d => formatter(d instanceof Date ? d : new Date(d))) as ValueFormatter<Value>
     }
 
-    return d3Format(format) as ValueFormatter
+    return (d3Format(format) as unknown) as ValueFormatter<Value>
 }
 
-export const computeGridLines = <Value>({
+export const computeGridLines = <Value extends string | number | Date>({
     width,
     height,
     scale,
diff --git a/packages/axes/src/types.ts b/packages/axes/src/types.ts
index c11c98e7d8..007ca3666e 100644
--- a/packages/axes/src/types.ts
+++ b/packages/axes/src/types.ts
@@ -10,8 +10,6 @@ import {
 } from 'd3-scale'
 import React from 'react'
 
-// export type AxisValue = string | number | Date
-
 export type GridValuesBuilder<T> = T extends number
     ? number[]
     : T extends string
@@ -39,7 +37,7 @@ export type AnyScale =
     | (ScaleLogarithmic<any, any> & { type: 'log' })
     | ScaleWithBandwidth
 
-export type TicksSpec<Value> =
+export type TicksSpec<Value extends string | number | Date> =
     // exact number of ticks, please note that
     // depending on the current range of values,
     // you might not get this exact count
@@ -53,46 +51,50 @@ export type TicksSpec<Value> =
 
 export type AxisLegendPosition = 'start' | 'middle' | 'end'
 
-export interface AxisProp<Value> {
+export type ValueFormatter<Value extends string | number | Date> = (value: Value) => Value
+
+export interface AxisProp<Value extends string | number | Date> {
     ticksPosition?: 'before' | 'after'
     tickValues?: TicksSpec<Value>
     tickSize?: number
     tickPadding?: number
     tickRotation?: number
-    format?: any
-    renderTick?: any
+    format?: string | ValueFormatter<Value>
+    renderTick?: (props: AxisTickProps<Value>) => JSX.Element
     legend?: React.ReactNode
     legendPosition?: AxisLegendPosition
     legendOffset?: number
 }
 
-export interface CanvasAxisProp<Value> extends Omit<AxisProp<Value>, 'legend'> {
+export interface CanvasAxisProp<Value extends string | number | Date>
+    extends Omit<AxisProp<Value>, 'legend'> {
     legend?: string
 }
 
-export interface AxisProps {
+export interface AxisProps<Value extends number | string | Date = any> {
     axis: 'x' | 'y'
     scale: AnyScale
     x?: number
     y?: number
     length: number
     ticksPosition: 'before' | 'after'
-    tickValues?: TicksSpec<number | string | Date>
+    tickValues?: TicksSpec<Value>
     tickSize?: number
     tickPadding?: number
     tickRotation?: number
-    format?: any
-    renderTick?: any
+    format?: string | ValueFormatter<Value>
+    renderTick?: (props: AxisTickProps<Value>) => JSX.Element
     legend?: React.ReactNode
     legendPosition?: 'start' | 'middle' | 'end'
     legendOffset?: number
-    onClick?: any
+    onClick?: (event: React.MouseEvent<SVGGElement, MouseEvent>, value: Value) => void
     ariaHidden?: boolean
 }
 
 export interface AxisTickProps<Value extends number | string | Date> {
+    tickIndex: number
     value: Value
-    format?: any
+    format?: string | ValueFormatter<Value>
     x: number
     y: number
     lineX: number
@@ -111,8 +113,6 @@ export interface AxisTickProps<Value extends number | string | Date> {
     onClick?: (event: React.MouseEvent<SVGGElement, MouseEvent>, value: Value) => void
 }
 
-export type ValueFormatter = (value: number | string) => string
-
 export type Line = {
     key: string
     x1: number