From 75f97bb1ffe96311105848d25dea119721213e08 Mon Sep 17 00:00:00 2001
From: Jordan <beauchjord@gmail.com>
Date: Mon, 30 Nov 2020 00:42:42 +1300
Subject: [PATCH] Allow typegen for custom RPCs (#2876)

* Allow typegen for custom RPCs

* It runs

* fix lint

* Remove unused metadata arg
---
 packages/typegen/src/fromChain.ts    |  5 +++--
 packages/typegen/src/generate/rpc.ts | 22 ++++++++++++++--------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/packages/typegen/src/fromChain.ts b/packages/typegen/src/fromChain.ts
index 5d9c9e622b..4966cf65ea 100644
--- a/packages/typegen/src/fromChain.ts
+++ b/packages/typegen/src/fromChain.ts
@@ -6,7 +6,7 @@ import yargs from 'yargs';
 import { formatNumber } from '@polkadot/util';
 import WS from '@polkadot/x-ws';
 
-import { generateDefaultConsts, generateDefaultQuery, generateDefaultTx } from './generate';
+import { generateDefaultConsts, generateDefaultQuery, generateDefaultRpc, generateDefaultTx } from './generate';
 import { HEADER, writeFile } from './util';
 
 function generate (metaHex: string, pkg: string | undefined, output: string, isStrict?: boolean): void {
@@ -19,6 +19,7 @@ function generate (metaHex: string, pkg: string | undefined, output: string, isS
 
   generateDefaultConsts(path.join(process.cwd(), output, 'augment-api-consts.ts'), metaHex, extraTypes, isStrict);
   generateDefaultQuery(path.join(process.cwd(), output, 'augment-api-query.ts'), metaHex, extraTypes, isStrict);
+  generateDefaultRpc(path.join(process.cwd(), output, 'augment-api-rpc.ts'), extraTypes);
   generateDefaultTx(path.join(process.cwd(), output, 'augment-api-tx.ts'), metaHex, extraTypes, isStrict);
 
   writeFile(path.join(process.cwd(), output, 'augment-api.ts'), (): string =>
@@ -26,7 +27,7 @@ function generate (metaHex: string, pkg: string | undefined, output: string, isS
       HEADER('chain'),
       ...[
         '@polkadot/api/augment/rpc',
-        ...['consts', 'query', 'tx'].filter((key) => !!key).map((key) => `./augment-api-${key}`)
+        ...['consts', 'query', 'tx', 'rpc'].filter((key) => !!key).map((key) => `./augment-api-${key}`)
       ].map((path) => `import '${path}';\n`)
     ].join('')
   );
diff --git a/packages/typegen/src/generate/rpc.ts b/packages/typegen/src/generate/rpc.ts
index a803e581c6..40e3951dda 100644
--- a/packages/typegen/src/generate/rpc.ts
+++ b/packages/typegen/src/generate/rpc.ts
@@ -7,7 +7,7 @@ import { TypeRegistry } from '@polkadot/types/create';
 import { Definitions } from '@polkadot/types/types';
 import * as defaultDefinitions from '@polkadot/types/interfaces/definitions';
 
-import { createImports, getSimilarTypes, formatType, readTemplate, setImports, writeFile } from '../util';
+import { createImports, getSimilarTypes, formatType, readTemplate, registerDefinitions, setImports, writeFile } from '../util';
 
 const StorageKeyTye = 'StorageKey | string | Uint8Array | any';
 
@@ -15,12 +15,12 @@ const template = readTemplate('rpc');
 const generateRpcTypesTemplate = Handlebars.compile(template);
 
 /** @internal */
-export function generateRpcTypes (importDefinitions: { [importPath: string]: Record<string, Definitions> }, dest: string): void {
+export function generateRpcTypes (registry: TypeRegistry, importDefinitions: Record<string, Definitions>, dest: string, extraTypes: Record<string, Record<string, { types: Record<string, any> }>> = {}): void {
   writeFile(dest, (): string => {
-    const registry = new TypeRegistry();
-    const imports = createImports(importDefinitions);
+    const allTypes: Record<string, Record<string, { types: Record<string, any> }>> = { '@polkadot/types/interfaces': importDefinitions, ...extraTypes };
+    const imports = createImports(allTypes);
     const definitions = imports.definitions as Record<string, Definitions>;
-    const allDefs = Object.entries(importDefinitions).reduce((defs, [path, obj]) => {
+    const allDefs = Object.entries(allTypes).reduce((defs, [path, obj]) => {
       return Object.entries(obj).reduce((defs, [key, value]) => ({ ...defs, [`${path}/${key}`]: value }), defs);
     }, {});
 
@@ -111,9 +111,15 @@ export function generateRpcTypes (importDefinitions: { [importPath: string]: Rec
   });
 }
 
-export function generateDefaultRpc (): void {
+export function generateDefaultRpc (dest = 'packages/api/src/augment/rpc.ts', extraTypes: Record<string, Record<string, { types: Record<string, any> }>> = {}): void {
+  const registry = new TypeRegistry();
+
+  registerDefinitions(registry, extraTypes);
+
   generateRpcTypes(
-    { '@polkadot/types/interfaces': defaultDefinitions },
-    'packages/api/src/augment/rpc.ts'
+    registry,
+    defaultDefinitions,
+    dest,
+    extraTypes
   );
 }