feat: add generic type arg to deployProxy, deployContract, deployBeac…

…onProxy, deployBeacon, upgradeProxy, upgradeBeacon, and forceImport functions

packages/plugin-hardhat/src/deploy-beacon-proxy.ts

@@ -24,27 +24,28 @@ import {
 } from './utils';
 import { enableDefender } from './defender/utils';
 import { getContractInstance } from './utils/contract-instance';
+import { ContractTypeOfFactory } from './type-extensions';
 
 export interface DeployBeaconProxyFunction {
-  (
+  <F extends ContractFactory>(
     beacon: ContractAddressOrInstance,
     attachTo: ContractFactory,
     args?: unknown[],
     opts?: DeployBeaconProxyOptions,
-  ): Promise<Contract>;
+  ): Promise<ContractTypeOfFactory<F>>;
   (beacon: ContractAddressOrInstance, attachTo: ContractFactory, opts?: DeployBeaconProxyOptions): Promise<Contract>;
 }
 
 export function makeDeployBeaconProxy(
   hre: HardhatRuntimeEnvironment,
   defenderModule: boolean,
 ): DeployBeaconProxyFunction {
-  return async function deployBeaconProxy(
+  return async function deployBeaconProxy<F extends ContractFactory>(
     beacon: ContractAddressOrInstance,
-    attachTo: ContractFactory,
+    attachTo: F,
     args: unknown[] | DeployBeaconProxyOptions = [],
     opts: DeployBeaconProxyOptions = {},
-  ) {
+  ): Promise<ContractTypeOfFactory<F>> {
     if (!(attachTo instanceof ContractFactory)) {
       throw new UpgradesError(
         `attachTo must specify a contract factory`,

packages/plugin-hardhat/src/deploy-beacon.ts

@@ -1,19 +1,23 @@
 import type { HardhatRuntimeEnvironment } from 'hardhat/types';
-import type { ContractFactory, Contract } from 'ethers';
+import type { ContractFactory } from 'ethers';
 
 import { Deployment } from '@openzeppelin/upgrades-core';
 
 import { DeployBeaconOptions, deploy, DeployTransaction, getUpgradeableBeaconFactory, deployBeaconImpl } from './utils';
 import { disableDefender } from './defender/utils';
 import { attach, getSigner } from './utils/ethers';
 import { getInitialOwner } from './utils/initial-owner';
+import { ContractTypeOfFactory } from './type-extensions';
 
 export interface DeployBeaconFunction {
-  (ImplFactory: ContractFactory, opts?: DeployBeaconOptions): Promise<Contract>;
+  <F extends ContractFactory>(ImplFactory: F, opts?: DeployBeaconOptions): Promise<ContractTypeOfFactory<F>>;
 }
 
 export function makeDeployBeacon(hre: HardhatRuntimeEnvironment, defenderModule: boolean): DeployBeaconFunction {
-  return async function deployBeacon(ImplFactory: ContractFactory, opts: DeployBeaconOptions = {}) {
+  return async function deployBeacon<F extends ContractFactory>(
+    ImplFactory: F,
+    opts: DeployBeaconOptions = {},
+  ): Promise<ContractTypeOfFactory<F>> {
     disableDefender(hre, defenderModule, opts, deployBeacon.name);
 
     const { impl } = await deployBeaconImpl(hre, ImplFactory, opts);
@@ -34,6 +38,6 @@ export function makeDeployBeacon(hre: HardhatRuntimeEnvironment, defenderModule:
 
     // @ts-ignore Won't be readonly because beaconContract was created through attach.
     beaconContract.deployTransaction = beaconDeployment.deployTransaction;
-    return beaconContract;
+    return beaconContract as ContractTypeOfFactory<F>;
   };
 }

packages/plugin-hardhat/src/deploy-contract.ts

@@ -1,5 +1,5 @@
 import { HardhatRuntimeEnvironment } from 'hardhat/types';
-import type { ContractFactory, Contract } from 'ethers';
+import type { ContractFactory } from 'ethers';
 
 import { deploy, DeployContractOptions, DeployTransaction } from './utils';
 import { DeployData, getDeployData } from './utils/deploy-impl';
@@ -13,17 +13,25 @@ import {
   inferInitializable,
 } from '@openzeppelin/upgrades-core';
 import { getContractInstance } from './utils/contract-instance';
+import { ContractTypeOfFactory } from './type-extensions';
 
 export interface DeployContractFunction {
-  (Contract: ContractFactory, args?: unknown[], opts?: DeployContractOptions): Promise<Contract>;
-  (Contract: ContractFactory, opts?: DeployContractOptions): Promise<Contract>;
+  <F extends ContractFactory>(
+    Contract: F,
+    args?: unknown[],
+    opts?: DeployContractOptions,
+  ): Promise<ContractTypeOfFactory<F>>;
+  <F extends ContractFactory>(
+    Contract: ContractFactory,
+    opts?: DeployContractOptions,
+  ): Promise<ContractTypeOfFactory<F>>;
 }
 
 async function deployNonUpgradeableContract(
   hre: HardhatRuntimeEnvironment,
   Contract: ContractFactory,
   opts: DeployContractOptions,
-) {
+): Promise<Deployment & DeployTransaction & RemoteDeploymentId> {
   const deployData = await getDeployData(hre, Contract, opts);
 
   if (!opts.unsafeAllowDeployContract) {
@@ -36,7 +44,6 @@ async function deployNonUpgradeableContract(
     Contract,
     ...deployData.fullOpts.constructorArgs,
   );
-
   return deployment;
 }
 
@@ -54,11 +61,11 @@ function assertNonUpgradeable(deployData: DeployData) {
 }
 
 export function makeDeployContract(hre: HardhatRuntimeEnvironment, defenderModule: boolean): DeployContractFunction {
-  return async function deployContract(
-    Contract,
+  return async function deployContract<F extends ContractFactory>(
+    Contract: F,
     args: unknown[] | DeployContractOptions = [],
     opts: DeployContractOptions = {},
-  ) {
+  ): Promise<ContractTypeOfFactory<F>> {
     if (!Array.isArray(args)) {
       opts = args;
       args = [];

packages/plugin-hardhat/src/deploy-proxy.ts

@@ -1,5 +1,5 @@
 import type { HardhatRuntimeEnvironment } from 'hardhat/types';
-import type { ContractFactory, Contract } from 'ethers';
+import type { ContractFactory } from 'ethers';
 
 import {
   Manifest,
@@ -25,18 +25,23 @@ import {
 import { enableDefender } from './defender/utils';
 import { getContractInstance } from './utils/contract-instance';
 import { getInitialOwner } from './utils/initial-owner';
+import { ContractTypeOfFactory } from './type-extensions';
 
 export interface DeployFunction {
-  (ImplFactory: ContractFactory, args?: unknown[], opts?: DeployProxyOptions): Promise<Contract>;
-  (ImplFactory: ContractFactory, opts?: DeployProxyOptions): Promise<Contract>;
+  <F extends ContractFactory>(
+    ImplFactory: F,
+    args?: unknown[],
+    opts?: DeployProxyOptions,
+  ): Promise<ContractTypeOfFactory<F>>;
+  <F extends ContractFactory>(ImplFactory: F, opts?: DeployProxyOptions): Promise<ContractTypeOfFactory<F>>;
 }
 
 export function makeDeployProxy(hre: HardhatRuntimeEnvironment, defenderModule: boolean): DeployFunction {
-  return async function deployProxy(
-    ImplFactory: ContractFactory,
+  return async function deployProxy<F extends ContractFactory>(
+    ImplFactory: F,
     args: unknown[] | DeployProxyOptions = [],
     opts: DeployProxyOptions = {},
-  ) {
+  ): Promise<ContractTypeOfFactory<F>> {
     if (!Array.isArray(args)) {
       opts = args;
       args = [];

packages/plugin-hardhat/src/force-import.ts

@@ -26,17 +26,22 @@ import {
 } from './utils';
 import { getDeployData } from './utils/deploy-impl';
 import { attach, getSigner } from './utils/ethers';
+import { ContractTypeOfFactory } from './type-extensions';
 
 export interface ForceImportFunction {
-  (proxyAddress: string, ImplFactory: ContractFactory, opts?: ForceImportOptions): Promise<Contract>;
+  <F extends ContractFactory>(
+    proxyAddress: string,
+    ImplFactory: F,
+    opts?: ForceImportOptions,
+  ): Promise<ContractTypeOfFactory<F> | Contract>;
 }
 
 export function makeForceImport(hre: HardhatRuntimeEnvironment): ForceImportFunction {
-  return async function forceImport(
+  return async function forceImport<F extends ContractFactory>(
     addressOrInstance: ContractAddressOrInstance,
-    ImplFactory: ContractFactory,
+    ImplFactory: F,
     opts: ForceImportOptions = {},
-  ) {
+  ): Promise<ContractTypeOfFactory<F> | Contract> {
     const { provider } = hre.network;
     const manifest = await Manifest.forNetwork(provider);
 
@@ -46,19 +51,19 @@ export function makeForceImport(hre: HardhatRuntimeEnvironment): ForceImportFunc
     if (implAddress !== undefined) {
       await importProxyToManifest(provider, hre, address, implAddress, ImplFactory, opts, manifest);
 
-      return attach(ImplFactory, address);
+      return attach(ImplFactory, address) as ContractTypeOfFactory<F>;
     } else if (await isBeacon(provider, address)) {
       const beaconImplAddress = await getImplementationAddressFromBeacon(provider, address);
       await addImplToManifest(hre, beaconImplAddress, ImplFactory, opts);
 
       const UpgradeableBeaconFactory = await getUpgradeableBeaconFactory(hre, getSigner(ImplFactory.runner));
-      return attach(UpgradeableBeaconFactory, address);
+      return attach(UpgradeableBeaconFactory, address) as Contract;
     } else {
       if (!(await hasCode(provider, address))) {
         throw new NoContractImportError(address);
       }
       await addImplToManifest(hre, address, ImplFactory, opts);
-      return attach(ImplFactory, address);
+      return attach(ImplFactory, address) as ContractTypeOfFactory<F>;
     }
   };
 }

packages/plugin-hardhat/src/type-extensions.ts

@@ -2,6 +2,9 @@ import 'hardhat/types/runtime';
 import 'hardhat/types/config';
 
 import type { HardhatUpgrades, DefenderHardhatUpgrades } from '.';
+import { ContractFactory } from 'ethers';
+
+export type ContractTypeOfFactory<F extends ContractFactory> = ReturnType<F['attach']>;
 
 declare module 'hardhat/types/runtime' {
   export interface HardhatRuntimeEnvironment {

packages/plugin-hardhat/src/upgrade-beacon.ts

@@ -1,5 +1,5 @@
 import { HardhatRuntimeEnvironment } from 'hardhat/types';
-import type { ContractFactory, Contract } from 'ethers';
+import type { ContractFactory } from 'ethers';
 
 import {
   getContractAddress,
@@ -11,15 +11,20 @@ import {
   getSigner,
 } from './utils';
 import { disableDefender } from './defender/utils';
+import { ContractTypeOfFactory } from './type-extensions';
 
-export type UpgradeBeaconFunction = (
+export type UpgradeBeaconFunction = <F extends ContractFactory>(
   beacon: ContractAddressOrInstance,
-  ImplFactory: ContractFactory,
+  ImplFactory: F,
   opts?: UpgradeBeaconOptions,
-) => Promise<Contract>;
+) => Promise<ContractTypeOfFactory<F>>;
 
 export function makeUpgradeBeacon(hre: HardhatRuntimeEnvironment, defenderModule: boolean): UpgradeBeaconFunction {
-  return async function upgradeBeacon(beacon, ImplFactory, opts: UpgradeBeaconOptions = {}) {
+  return async function upgradeBeacon<F extends ContractFactory>(
+    beacon: ContractAddressOrInstance,
+    ImplFactory: F,
+    opts: UpgradeBeaconOptions = {},
+  ): Promise<ContractTypeOfFactory<F>> {
     disableDefender(hre, defenderModule, opts, upgradeBeacon.name);
 
     const beaconAddress = await getContractAddress(beacon);
@@ -33,6 +38,6 @@ export function makeUpgradeBeacon(hre: HardhatRuntimeEnvironment, defenderModule
 
     // @ts-ignore Won't be readonly because beaconContract was created through attach.
     beaconContract.deployTransaction = upgradeTx;
-    return beaconContract;
+    return beaconContract as ContractTypeOfFactory<F>;
   };
 }

packages/plugin-hardhat/src/upgrade-proxy.ts

@@ -1,5 +1,5 @@
 import { HardhatRuntimeEnvironment } from 'hardhat/types';
-import type { ethers, ContractFactory, Contract, Signer } from 'ethers';
+import type { ethers, ContractFactory, Signer } from 'ethers';
 import debug from './utils/debug';
 import { getAdminAddress, getCode, getUpgradeInterfaceVersion, isEmptySlot } from '@openzeppelin/upgrades-core';
 
@@ -18,19 +18,24 @@ import {
   attachProxyAdminV4,
   attachProxyAdminV5,
 } from './utils/attach-abi';
+import { ContractTypeOfFactory } from './type-extensions';
 
-export type UpgradeFunction = (
+export type UpgradeFunction = <F extends ContractFactory>(
   proxy: ContractAddressOrInstance,
-  ImplFactory: ContractFactory,
+  ImplFactory: F,
   opts?: UpgradeProxyOptions,
-) => Promise<Contract>;
+) => Promise<ContractTypeOfFactory<F>>;
 
 export function makeUpgradeProxy(
   hre: HardhatRuntimeEnvironment,
   defenderModule: boolean,
   log = debug,
 ): UpgradeFunction {
-  return async function upgradeProxy(proxy, ImplFactory, opts: UpgradeProxyOptions = {}) {
+  return async function upgradeProxy<F extends ContractFactory>(
+    proxy: ContractAddressOrInstance,
+    ImplFactory: F,
+    opts: UpgradeProxyOptions = {},
+  ): Promise<ContractTypeOfFactory<F>> {
     disableDefender(hre, defenderModule, opts, upgradeProxy.name);
 
     const proxyAddress = await getContractAddress(proxy);
@@ -44,7 +49,7 @@ export function makeUpgradeProxy(
     const inst = attach(ImplFactory, proxyAddress);
     // @ts-ignore Won't be readonly because inst was created through attach.
     inst.deployTransaction = upgradeTx;
-    return inst;
+    return inst as ContractTypeOfFactory<F>;
   };
 
   type Upgrader = (nextImpl: string, call?: string) => Promise<ethers.TransactionResponse>;

packages/plugin-hardhat/src/utils/contract-instance.ts

@@ -6,6 +6,7 @@ import { DeployTransaction, DefenderDeploy } from '.';
 import { waitForDeployment } from '../defender/utils';
 import { Deployment, RemoteDeploymentId, DeployOpts } from '@openzeppelin/upgrades-core';
 import { attach } from './ethers';
+import { ContractTypeOfFactory } from '../type-extensions';
 
 /**
  * Gets a contract instance from a deployment, where the deployment may be remote.
@@ -19,13 +20,13 @@ import { attach } from './ethers';
  * @param deployTransaction The transaction that deployed the contract, if available
  * @returns The contract instance
  */
-export function getContractInstance(
+export function getContractInstance<F extends ContractFactory>(
   hre: HardhatRuntimeEnvironment,
-  contract: ContractFactory,
+  contract: F,
   opts: DeployOpts & DefenderDeploy,
   deployment: Deployment & DeployTransaction & RemoteDeploymentId,
-) {
-  const instance = attach(contract, deployment.address);
+): ContractTypeOfFactory<F> {
+  const instance = attach(contract, deployment.address) as ContractTypeOfFactory<F>;
 
   // @ts-ignore Won't be readonly because instance was created through attach.
   instance.deploymentTransaction = () => deployment.deployTransaction ?? null; // Convert undefined to null to conform to ethers.js types.

packages/plugin-hardhat/src/utils/ethers.ts

@@ -1,4 +1,4 @@
-import { Contract, ContractFactory, ContractRunner, Signer } from 'ethers';
+import { ContractFactory, ContractRunner, Signer, Contract } from 'ethers';
 
 /**
  * Attaches a ContractFactory to an address and returns a Contract instance.