Merge pull request argentlabs#1971 from argentlabs/feat/zod-token-schema

feat: use zod token schema

packages/extension/src/shared/token/storage.ts

@@ -1,9 +1,7 @@
-import * as yup from "yup"
+import { z } from "zod"
 
-import { addressSchema } from "../../ui/services/addresses"
 import { ArrayStorage } from "../storage"
-import { assertSchema } from "../utils/schema"
-import { BaseToken, Token } from "./type"
+import { BaseToken, BaseTokenSchema, Token } from "./type"
 import { equalToken, parsedDefaultTokens } from "./utils"
 
 export const tokenStore = new ArrayStorage(parsedDefaultTokens, {
@@ -12,40 +10,34 @@ export const tokenStore = new ArrayStorage(parsedDefaultTokens, {
   compare: equalToken,
 })
 
-export const baseTokenSchema: yup.Schema<BaseToken> = yup
-  .object()
-  .required("BaseToken is required")
-  .shape({
-    address: addressSchema.required("Address is required"),
-    networkId: yup.string().required("Network is required"),
-  })
+export const tokenSchema = BaseTokenSchema.extend({
+  name: z.string({ required_error: "Name is required" }),
+  symbol: z
+    .string({
+      required_error: "Symbol is required",
+    })
+    .min(1, { message: "Symbol must be atleast 1 character" }),
+  decimals: z.number({ required_error: "Decimals is required" }),
+  image: z.string().optional(),
+  showAlways: z.boolean().optional(),
+})
 
-export const tokenSchema: yup.Schema<Token> = baseTokenSchema
-  .required("Token is required")
-  .shape({
-    name: yup.string().required("Name is required"),
-    symbol: yup.string().required("Symbol is required").min(1).max(6),
-    decimals: yup
-      .string()
-      .matches(/^\d+$/, "Decimals must be a number")
-      .required("Decimals is required"),
-    image: yup.string(),
-    showAlways: yup.boolean(),
+export async function addToken(token: Token) {
+  const newToken: Token = tokenSchema.parse({
+    ...token,
+    showAlways: true,
   })
 
-export async function addToken(token: Token) {
-  const newToken: Token = { ...token, showAlways: true }
-  await assertSchema(tokenSchema, newToken)
   return tokenStore.push(newToken)
 }
 
 export async function hasToken(token: BaseToken) {
-  await assertSchema(baseTokenSchema, token)
-  const [hit] = await tokenStore.get((t) => equalToken(t, token))
+  const parsedToken = BaseTokenSchema.parse(token)
+  const [hit] = await tokenStore.get((t) => equalToken(t, parsedToken))
   return Boolean(hit)
 }
 
 export async function removeToken(token: BaseToken) {
-  await assertSchema(baseTokenSchema, token)
-  return tokenStore.remove((t) => equalToken(t, token))
+  const parsedToken = BaseTokenSchema.parse(token)
+  return tokenStore.remove((t) => equalToken(t, parsedToken))
 }

packages/extension/src/shared/token/type.ts

@@ -1,17 +1,28 @@
-export interface BaseToken {
-  address: string
-  networkId: string
-}
+import { z } from "zod"
 
-export interface RequestToken extends Omit<BaseToken, "networkId"> {
-  address: string
-  networkId?: string
-  name?: string
-  symbol?: string
-  decimals?: number
-}
+export const BaseTokenSchema = z.object(
+  {
+    address: z.string({ required_error: "Address is required" }),
+    networkId: z.string({ required_error: "Network is required" }),
+  },
+  { required_error: "BaseToken is required" },
+)
 
-export interface Token extends Required<RequestToken> {
-  image?: string
-  showAlways?: boolean
-}
+export type BaseToken = z.infer<typeof BaseTokenSchema>
+
+export const RequestTokenSchema = z.object({
+  address: z.string(),
+  networkId: z.string().optional(),
+  name: z.string().optional(),
+  symbol: z.string().optional(),
+  decimals: z.number().optional(),
+})
+
+export type RequestToken = z.infer<typeof RequestTokenSchema>
+
+export const TokenSchema = RequestTokenSchema.required().extend({
+  image: z.string().optional(),
+  showAlways: z.boolean().optional(),
+})
+
+export type Token = z.infer<typeof TokenSchema>

packages/extension/src/ui/features/actions/AddTokenScreen.tsx

@@ -4,6 +4,7 @@ import React, { FC, useEffect, useMemo, useRef, useState } from "react"
 import { useNavigate } from "react-router-dom"
 import { number } from "starknet"
 import styled from "styled-components"
+import { ZodError } from "zod"
 
 import { addToken } from "../../../shared/token/storage"
 import { RequestToken, Token } from "../../../shared/token/type"
@@ -181,7 +182,11 @@ export const AddTokenScreen: FC<AddTokenScreenProps> = ({
                 onSubmit?.()
                 navigate(routes.accountTokens())
               } catch (e) {
-                setError("Token already exists")
+                if (e instanceof ZodError) {
+                  setError(e.issues[0].message)
+                } else {
+                  setError("Token not supported")
+                }
               }
             }
           }}