Refactor

packages/core/src/contract.ts

@@ -5,7 +5,7 @@ export interface Contract {
   license: string;
   parents: Parent[];
   natspecTags: NatspecTag[];
-  imports: ParentContract[];
+  imports: ImportContract[];
   functions: ContractFunction[];
   constructorCode: string[];
   constructorArgs: FunctionArgument[];
@@ -16,11 +16,11 @@ export interface Contract {
 export type Value = string | number | { lit: string } | { note: string, value: Value };
 
 export interface Parent {
-  contract: ParentContract;
+  contract: ImportContract;
   params: Value[];
 }
 
-export interface ParentContract extends ReferencedContract {
+export interface ImportContract extends ReferencedContract {
   path: string;
 }
 
@@ -29,7 +29,7 @@ export interface ReferencedContract {
 }
 
 export interface Using {
-  library: ParentContract;
+  library: ImportContract;
   usingFor: string;
 }
 
@@ -103,7 +103,7 @@ export class ContractBuilder implements Contract {
     });
   }
 
-  get imports(): ParentContract[] {
+  get imports(): ImportContract[] {
     return [...this.importFromMap.entries()].map(([name, path]) => ({ name, path }));
   }
 
@@ -115,15 +115,15 @@ export class ContractBuilder implements Contract {
     return [...this.variableSet];
   }
 
-  addParent(contract: ParentContract, params: Value[] = []): boolean {
+  addParent(contract: ImportContract, params: Value[] = []): boolean {
     const present = this.parentMap.has(contract.name);
     this.parentMap.set(contract.name, { contract, params });
-    this.addImportFrom(contract.name, contract.path);
+    this.addImportFrom(contract);
     return !present;
   }
 
-  addImportFrom(importName: string, from: string) {
-    this.importFromMap.set(importName, from);
+  addImportFrom(contract: ImportContract) {
+    this.importFromMap.set(contract.name, contract.path);
   }
 
   addOverride(parent: ReferencedContract, baseFn: BaseFunction, mutability?: FunctionMutability) {

packages/core/src/erc20.ts

@@ -184,7 +184,10 @@ function addVotes(c: ContractBuilder, clockMode: ClockMode) {
   c.addParent(ERC20Votes);
   c.addOverride(ERC20Votes, functions._update);
 
-  c.addImportFrom('Nonces', '@openzeppelin/contracts/utils/Nonces.sol');
+  c.addImportFrom({
+    name: 'Nonces',
+    path: '@openzeppelin/contracts/utils/Nonces.sol',
+  });
   c.addOverride({
     name: 'Nonces',
   }, functions.nonces);

packages/core/src/governor.ts

@@ -250,7 +250,10 @@ function addCounting(c: ContractBuilder) {
 function addVotes(c: ContractBuilder) {
   const tokenArg = '_token';
 
-  c.addImportFrom('IVotes', `@openzeppelin/contracts/governance/extensions/GovernorVotes.sol`);
+  c.addImportFrom({
+    name: 'IVotes',
+    path: `@openzeppelin/contracts/governance/extensions/GovernorVotes.sol`,
+  });
   c.addConstructorArgument({
     type: {
       name: 'IVotes',
@@ -360,7 +363,7 @@ function addTimelock(c: ContractBuilder, { timelock }: Required<GovernorOptions>
   const timelockArg = '_timelock';
   const { timelockType, timelockParent } = timelockModules[timelock];
 
-  c.addImportFrom(timelockType.name, timelockType.path);
+  c.addImportFrom(timelockType);
   c.addConstructorArgument({
     type: timelockType,
     name: timelockArg,

packages/core/src/options.ts

@@ -1,6 +1,6 @@
 import path from 'path';
 
-import type { Contract, ReferencedContract, ParentContract } from './contract';
+import type { Contract, ReferencedContract, ImportContract } from './contract';
 import { inferTranspiled } from './infer-transpiled';
 
 const upgradeableName = (n: string) => {
@@ -11,7 +11,7 @@ const upgradeableName = (n: string) => {
   }
 }
 
-const upgradeableImport = (p: ParentContract): ParentContract => {
+const upgradeableImport = (p: ImportContract): ImportContract => {
   const { dir, ext, name } = path.parse(p.path);
   // Use path.posix to get forward slashes
   return {
@@ -26,7 +26,7 @@ const upgradeableImport = (p: ParentContract): ParentContract => {
 };
 
 export interface Options {
-  transformImport?: (parent: ParentContract) => ParentContract;
+  transformImport?: (parent: ImportContract) => ImportContract;
 }
 
 export interface Helpers extends Required<Options> {

packages/core/src/print.ts

@@ -1,6 +1,6 @@
 import 'array.prototype.flatmap/auto';
 
-import type { Contract, Parent, ContractFunction, FunctionArgument, Value, NatspecTag, ParentContract } from './contract';
+import type { Contract, Parent, ContractFunction, FunctionArgument, Value, NatspecTag, ImportContract } from './contract';
 import { Options, Helpers, withHelpers } from './options';
 
 import { formatLines, spaceBetween, Lines } from './utils/format-lines';
@@ -247,9 +247,9 @@ function printNatspecTags(tags: NatspecTag[]): string[] {
   return tags.map(({ key, value }) => `/// ${key} ${value}`);
 }
 
-function printImports(imports: ParentContract[], helpers: Helpers): string[] {
+function printImports(imports: ImportContract[], helpers: Helpers): string[] {
   // Combine imports by path
-  const pathMap = new Map<string, ParentContract[]>();
+  const pathMap = new Map<string, ImportContract[]>();
   for (const imp of imports) {
     const path = imp.path;
     if (!pathMap.has(path)) {

packages/core/src/stablecoin.ts

@@ -235,7 +235,10 @@ function addCustodian(c: ContractBuilder, access: Access) {
       break;
     }
     case 'managed': {
-      c.addImportFrom('AuthorityUtils', `@openzeppelin/contracts/access/manager/AuthorityUtils.sol`);
+      c.addImportFrom({
+        name: 'AuthorityUtils',
+        path: `@openzeppelin/contracts/access/manager/AuthorityUtils.sol`,
+      });
       const logic = [
         `(bool immediate,) = AuthorityUtils.canCallWithDelay(authority(), user, address(this), bytes4(_msgData()[0:4]));`,
         `return immediate;`
@@ -268,7 +271,10 @@ function addVotes(c: ContractBuilder, clockMode: ClockMode) {
   c.addParent(ERC20Votes);
   c.addOverride(ERC20Votes, functions._update);
 
-  c.addImportFrom('Nonces', '@openzeppelin/contracts/utils/Nonces.sol');
+  c.addImportFrom({
+    name: 'Nonces',
+    path: '@openzeppelin/contracts/utils/Nonces.sol',
+  });
   c.addOverride({
     name: 'Nonces',
   }, functions.nonces);