codegen: Move DataSourceTemplate logic to protocols folder

Author: evaporei

src/codegen/template.js

@@ -3,14 +3,19 @@ const immutable = require('immutable')
 const tsCodegen = require('./typescript')
 
 module.exports = class DataSourceTemplateCodeGenerator {
-  constructor(template) {
+  constructor(template, protocol) {
     this.template = template
+    this.protocolTemplateCodeGen = protocol.getTemplateCodeGen(template)
   }
 
   generateModuleImports() {
     return [
       tsCodegen.moduleImports(
-        ['Address', 'DataSourceTemplate', 'DataSourceContext'],
+        [
+          ...this.protocolTemplateCodeGen.generateModuleImports(),
+          'DataSourceTemplate',
+          'DataSourceContext',
+        ],
         '@graphprotocol/graph-ts',
       ),
     ]
@@ -24,55 +29,8 @@ module.exports = class DataSourceTemplateCodeGenerator {
     let name = this.template.get('name')
 
     let klass = tsCodegen.klass(name, { export: true, extends: 'DataSourceTemplate' })
-    klass.addMethod(this._generateCreateMethod())
-    klass.addMethod(this._generateCreateWithContextMethod())
+    klass.addMethod(this.protocolTemplateCodeGen.generateCreateMethod())
+    klass.addMethod(this.protocolTemplateCodeGen.generateCreateWithContextMethod())
     return klass
   }
-
-  _generateCreateMethod() {
-    let name = this.template.get('name')
-    let kind = this.template.get('kind')
-
-    switch (kind) {
-      case 'ethereum/contract':
-        return tsCodegen.staticMethod(
-          'create',
-          [tsCodegen.param('address', tsCodegen.namedType('Address'))],
-          tsCodegen.namedType('void'),
-          `
-          DataSourceTemplate.create('${name}', [address.toHex()])
-          `,
-        )
-
-      default:
-        throw new Error(
-          `Data sources with kind != 'ethereum/contract' are not supported yet`,
-        )
-    }
-  }
-
-  _generateCreateWithContextMethod() {
-    let name = this.template.get('name')
-    let kind = this.template.get('kind')
-
-    switch (kind) {
-      case 'ethereum/contract':
-        return tsCodegen.staticMethod(
-          'createWithContext',
-          [
-            tsCodegen.param('address', tsCodegen.namedType('Address')),
-            tsCodegen.param('context', tsCodegen.namedType('DataSourceContext')),
-          ],
-          tsCodegen.namedType('void'),
-          `
-          DataSourceTemplate.createWithContext('${name}', [address.toHex()], context)
-          `,
-        )
-
-      default:
-        throw new Error(
-          `Data sources with kind != 'ethereum/contract' are not supported yet`,
-        )
-    }
-  }
 }

src/protocols/ethereum/codegen/template.js

@@ -0,0 +1,42 @@
+const tsCodegen = require('../../../codegen/typescript')
+
+module.exports = class EthereumTemplateCodeGen {
+  constructor(template) {
+    this.template = template
+  }
+
+  generateModuleImports() {
+    return [
+      'Address',
+    ]
+  }
+
+  generateCreateMethod() {
+    const name = this.template.get('name')
+
+    return tsCodegen.staticMethod(
+      'create',
+      [tsCodegen.param('address', tsCodegen.namedType('Address'))],
+      tsCodegen.namedType('void'),
+      `
+      DataSourceTemplate.create('${name}', [address.toHex()])
+      `,
+    )
+  }
+
+  generateCreateWithContextMethod() {
+    const name = this.template.get('name')
+
+    return tsCodegen.staticMethod(
+      'createWithContext',
+      [
+        tsCodegen.param('address', tsCodegen.namedType('Address')),
+        tsCodegen.param('context', tsCodegen.namedType('DataSourceContext')),
+      ],
+      tsCodegen.namedType('void'),
+      `
+      DataSourceTemplate.createWithContext('${name}', [address.toHex()], context)
+      `,
+    )
+  }
+}

src/protocols/index.js

@@ -1,4 +1,6 @@
 const EthereumTypeGenerator = require('./ethereum/type-generator')
+const EthereumTemplateCodeGen = require('./ethereum/codegen/template')
+const NearTemplateCodeGen = require('./near/codegen/template')
 
 module.exports = class Protocol {
   static fromDataSources(dataSourcesAndTemplates) {
@@ -29,4 +31,18 @@ module.exports = class Protocol {
         return null
     }
   }
+
+  getTemplateCodeGen(template) {
+    switch (this.name) {
+      case 'ethereum':
+      case 'ethereum/contract':
+        return new EthereumTemplateCodeGen(template)
+      case 'near':
+        return new NearTemplateCodeGen(template)
+      default:
+        throw new Error(
+          `Data sources with kind '${this.name}' are not supported yet`,
+        )
+    }
+  }
 }

src/protocols/near/codegen/template.js

@@ -0,0 +1,42 @@
+const tsCodegen = require('../../../codegen/typescript')
+
+module.exports = class NearTemplateCodeGen {
+  constructor(template) {
+    this.template = template
+  }
+
+  generateModuleImports() {
+    return [
+      'near',
+    ]
+  }
+
+  generateCreateMethod() {
+    const name = this.template.get('name')
+
+    return tsCodegen.staticMethod(
+      'create',
+      [tsCodegen.param('accountId', tsCodegen.namedType('near.AccountId'))],
+      tsCodegen.namedType('void'),
+      `
+      DataSourceTemplate.create('${name}', [accountId])
+      `,
+    )
+  }
+
+  generateCreateWithContextMethod() {
+    const name = this.template.get('name')
+
+    return tsCodegen.staticMethod(
+      'createWithContext',
+      [
+        tsCodegen.param('accountId', tsCodegen.namedType('near.AccountId')),
+        tsCodegen.param('context', tsCodegen.namedType('DataSourceContext')),
+      ],
+      tsCodegen.namedType('void'),
+      `
+      DataSourceTemplate.createWithContext('${name}', [accountId], context)
+      `,
+    )
+  }
+}

src/type-generator.js

@@ -156,7 +156,7 @@ module.exports = class TypeGenerator {
               `${template.get('name')}`,
             )
 
-            let codeGenerator = new DataSourceTemplateCodeGenerator(template)
+            let codeGenerator = new DataSourceTemplateCodeGenerator(template, this.protocol)
 
             // Only generate module imports once, because they are identical for
             // all types generated for data source templates.