feat: add more strict types and better inference

src/TsAspectContainer.ts

@@ -1,11 +1,11 @@
 import { Advice } from './advice.enum';
 import { Aspect } from './aspect.interface';
 
-type AdviceAspectMap = Map<Advice, Aspect[]>;
-type MethodContainer = {
+type AdviceAspectMap<Data, Args> = Map<Advice, Aspect<Data, Args>[]>;
+type MethodContainer<Data, Args> = {
     originalMethod: any;
-    adviceAspectMap: AdviceAspectMap;
+    adviceAspectMap: AdviceAspectMap<Data, Args>;
 };
-type TsAspectContainer = Record<string, MethodContainer>;
+type TsAspectContainer = Record<string, MethodContainer<any, any>>;
 
 export { TsAspectContainer, MethodContainer, AdviceAspectMap };

src/addAspect.ts

@@ -4,7 +4,12 @@ import { Aspect } from './aspect.interface';
 import { proxyFunc, asyncProxyFunc } from './proxyFunc';
 import { setTsAspectProp, getTsAspectProp } from './TsAspectProperty';
 
-export function addAspect(target: any, methodName: string, advice: Advice, aspect: Aspect): void {
+export function addAspect<Data, Args>(
+    target: any,
+    methodName: string,
+    advice: Advice,
+    aspect: Aspect<Data, Args>,
+): void {
     let tsAspectProp = getTsAspectProp(target);
     if (!tsAspectProp) {
         tsAspectProp = {};
@@ -16,7 +21,7 @@ export function addAspect(target: any, methodName: string, advice: Advice, aspec
 
         tsAspectProp[methodName] = {
             originalMethod,
-            adviceAspectMap: new Map<Advice, Aspect[]>(),
+            adviceAspectMap: new Map<Advice, Aspect<Data, Args>[]>(),
         };
 
         const wrapperFunc = function (...args: any): any {

src/addAspectToPointcut.ts

@@ -3,11 +3,11 @@ import { Advice } from './advice.enum';
 import { Aspect } from './aspect.interface';
 import { getPointcutMethods } from './getPointcutMethods';
 
-export function addAspectToPointcut(
+export function addAspectToPointcut<Data, Args>(
     target: any,
     pointcut: string,
     advice: Advice,
-    aspect: Aspect,
+    aspect: Aspect<Data, Args>,
 ): void {
     const methods = getPointcutMethods(target, pointcut);
     methods.forEach(method => {

src/aspect.interface.ts

@@ -1,11 +1,11 @@
-export interface Aspect {
-    execute(ctx: AspectContext): any;
+export interface Aspect<Data = any, Args = any> {
+    execute(ctx: AspectContext<Data, Args>): any;
 }
 
-export type AspectContext = {
+export type AspectContext<Data, Args> = {
     target: any;
     methodName: string;
-    functionParams: any[];
-    returnValue: any;
+    functionParams: Args;
+    returnValue: Data | null;
     error: any;
 };

src/decorator/UseAspect.ts

@@ -1,11 +1,16 @@
 import { types } from 'util';
+import { getTsAspectProp, setTsAspectProp } from '../TsAspectProperty';
 import { Advice } from '../advice.enum';
 import { Aspect } from '../aspect.interface';
-import { proxyFunc, asyncProxyFunc } from '../proxyFunc';
-import { getTsAspectProp, setTsAspectProp } from '../TsAspectProperty';
+import { asyncProxyFunc, proxyFunc } from '../proxyFunc';
 
-export function UseAspect(advice: Advice, aspect: Aspect | (new () => Aspect)): MethodDecorator {
-    return function (target, propertyKey: string | symbol, descriptor: PropertyDescriptor) {
+export function UseAspect<T extends (...args: any[]) => any>(
+    advice: Advice,
+    aspect:
+        | Aspect<Awaited<ReturnType<T>>, Parameters<T>>
+        | (new () => Aspect<Awaited<ReturnType<T>>, Parameters<T>>),
+) {
+    return (target: any, propertyKey: string | symbol, descriptor: TypedPropertyDescriptor<T>) => {
         let tsAspectProp = getTsAspectProp(target);
         if (!tsAspectProp) {
             tsAspectProp = {};
@@ -19,7 +24,7 @@ export function UseAspect(advice: Advice, aspect: Aspect | (new () => Aspect)):
 
             tsAspectProp[propertyKeyString] = {
                 originalMethod,
-                adviceAspectMap: new Map<Advice, Aspect[]>(),
+                adviceAspectMap: new Map<Advice, Aspect<Awaited<ReturnType<T>>, Parameters<T>>[]>(),
             };
 
             descriptor.value = function (...args: any): any {
@@ -41,8 +46,8 @@ export function UseAspect(advice: Advice, aspect: Aspect | (new () => Aspect)):
                         );
                     }
                 }
-                return originalMethod(...args);
-            };
+                return originalMethod?.(...args);
+            } as T;
         }
 
         const { adviceAspectMap } = tsAspectProp[propertyKeyString];

src/proxyFunc.ts

@@ -2,14 +2,14 @@ import { Advice } from './advice.enum';
 import { AspectContext } from './aspect.interface';
 import { AdviceAspectMap, MethodContainer } from './TsAspectContainer';
 
-export async function asyncProxyFunc(
+export async function asyncProxyFunc<Data, Args>(
     target: any,
     methodName: string,
-    methodContainer: MethodContainer,
+    methodContainer: MethodContainer<Data, Args>,
     ...args: any
 ): Promise<any> {
     const { originalMethod, adviceAspectMap } = methodContainer;
-    const aspectCtx: AspectContext = {
+    const aspectCtx: AspectContext<Data, Args> = {
         target: target,
         methodName: methodName,
         functionParams: args,
@@ -43,14 +43,14 @@ export async function asyncProxyFunc(
     return aspectCtx.returnValue;
 }
 
-export function proxyFunc(
+export function proxyFunc<Data, Args>(
     target: any,
     methodName: string,
-    methodContainer: MethodContainer,
+    methodContainer: MethodContainer<Data, Args>,
     ...args: any
 ): any {
     const { originalMethod, adviceAspectMap } = methodContainer;
-    const aspectCtx: AspectContext = {
+    const aspectCtx: AspectContext<Data, Args> = {
         target: target,
         methodName: methodName,
         functionParams: args,
@@ -84,9 +84,9 @@ export function proxyFunc(
     return aspectCtx.returnValue;
 }
 
-function applyPreExecutionAspects(
-    aspectCtx: AspectContext,
-    adviceAspectMap: AdviceAspectMap,
+function applyPreExecutionAspects<Data, Args>(
+    aspectCtx: AspectContext<Data, Args>,
+    adviceAspectMap: AdviceAspectMap<Data, Args>,
 ): void {
     if (adviceAspectMap.has(Advice.Before)) {
         adviceAspectMap.get(Advice.Before)?.forEach(aspect => {
@@ -100,9 +100,9 @@ function applyPreExecutionAspects(
     }
 }
 
-function applyPostExecutionAspects(
-    aspectCtx: AspectContext,
-    adviceAspectMap: AdviceAspectMap,
+function applyPostExecutionAspects<Data, Args>(
+    aspectCtx: AspectContext<Data, Args>,
+    adviceAspectMap: AdviceAspectMap<Data, Args>,
 ): void {
     if (adviceAspectMap.has(Advice.Around)) {
         adviceAspectMap.get(Advice.Around)?.forEach(aspect => {

test/addAspect.test.ts

@@ -7,7 +7,7 @@ import { CalculatorCls } from './samples/CalculatorCls.sample';
 
 describe('addAspect', () => {
     let calculator: CalculatorCls;
-    const aspect = mock<Aspect>();
+    const aspect = mock<Aspect<number, [number, number]>>();
 
     beforeEach(() => {
         jest.clearAllMocks();
@@ -63,7 +63,7 @@ describe('addAspect', () => {
 
             calculator.add(1, 2);
 
-            const expectedCtx: AspectContext = {
+            const expectedCtx: AspectContext<number, [number, number]> = {
                 target: calculator,
                 methodName: 'add',
                 functionParams: [1, 2],
@@ -81,7 +81,7 @@ describe('addAspect', () => {
 
         expect(aspect.execute).toHaveBeenCalledTimes(1);
 
-        const expectedCtx: AspectContext = {
+        const expectedCtx: AspectContext<number, [number, number]> = {
             target: calculator,
             methodName: 'divide',
             functionParams: [1, 0],
@@ -92,7 +92,7 @@ describe('addAspect', () => {
     });
 
     it('should pass the returned value to the injected aspect for Advice.AfterReturn', () => {
-        aspect.execute.mockImplementationOnce((ctx: AspectContext) => {
+        aspect.execute.mockImplementationOnce((ctx: AspectContext<number, [number, number]>) => {
             return ctx.returnValue;
         });
 
@@ -102,7 +102,7 @@ describe('addAspect', () => {
 
         expect(aspect.execute).toHaveBeenCalledTimes(1);
 
-        const expectedCtx: AspectContext = {
+        const expectedCtx: AspectContext<number, [number, number]> = {
             target: calculator,
             methodName: 'add',
             functionParams: [1, 2],
@@ -113,8 +113,11 @@ describe('addAspect', () => {
     });
 
     it('should return the returned value manipulated by the injected aspect for Advice.AfterReturn', () => {
-        aspect.execute.mockImplementationOnce((ctx: AspectContext) => {
+        aspect.execute.mockImplementationOnce((ctx: AspectContext<number, [number, number]>) => {
             const returnValue = ctx.returnValue;
+            if (returnValue === null) {
+                return null;
+            }
             return returnValue * 42;
         });
 
@@ -126,8 +129,8 @@ describe('addAspect', () => {
     });
 
     it('should execute all injected aspects for the same advice and pointcut', () => {
-        const secondAspect = mock<Aspect>();
-        const thirdAspect = mock<Aspect>();
+        const secondAspect = mock<Aspect<number, [number, number]>>();
+        const thirdAspect = mock<Aspect<number, [number, number]>>();
 
         addAspect(calculator, 'add', Advice.Before, aspect);
 
@@ -160,8 +163,8 @@ describe('addAspect', () => {
             }
         }
 
-        const aspect = mock<Aspect>();
-        aspect.execute.mockImplementation((ctx: AspectContext) => {
+        const aspect = mock<Aspect<number, [number, number]>>();
+        aspect.execute.mockImplementation((ctx: AspectContext<number, [number, number]>) => {
             return 42;
         });
 

test/addAspectToPointcut.test.ts

@@ -7,7 +7,7 @@ import { CalculatorCls } from './samples/CalculatorCls.sample';
 
 describe('addAspectToPointcut', () => {
     let calculator: CalculatorCls;
-    const aspect = mock<Aspect>();
+    const aspect = mock<Aspect<number, [number, number]>>();
 
     beforeEach(() => {
         jest.clearAllMocks();

test/decorator/UseAspect.test.ts

@@ -4,8 +4,17 @@ import { Advice } from '../../src/advice.enum';
 import { Aspect, AspectContext } from '../../src/aspect.interface';
 import { UseAspect } from '../../src/decorator/UseAspect';
 
-const beforeAspect = mock<Aspect>();
-const afterAspect = mock<Aspect>();
+const beforeAspect = mock<Aspect<number, []>>();
+const afterAspect = mock<Aspect<void, [number]>>();
+
+const throwBeforeAspect = mock<Aspect<void, []>>();
+const throwAfterAspect = mock<Aspect<void, []>>();
+
+class TestAspect implements Aspect {
+    execute(ctx: AspectContext<void, [number]>) {
+        throw new Error('Method not implemented.');
+    }
+}
 
 class SampleClass {
     public constructor(private sampleId: number) {}
@@ -22,8 +31,8 @@ class SampleClass {
         this.sampleId = sampleId;
     }
 
-    @UseAspect(Advice.Before, beforeAspect)
-    @UseAspect(Advice.After, afterAspect)
+    @UseAspect(Advice.Before, throwBeforeAspect)
+    @UseAspect(Advice.After, throwAfterAspect)
     public throwError(): void {
         throw new Error('this is expected!');
     }
@@ -43,7 +52,7 @@ describe('UseAspect', () => {
 
         expect(beforeAspect.execute).toHaveBeenCalledTimes(1);
 
-        const expectedCtx: AspectContext = {
+        const expectedCtx: AspectContext<number, []> = {
             target: sample,
             methodName: 'getSampleId',
             functionParams: [],
@@ -55,8 +64,8 @@ describe('UseAspect', () => {
 
     it('should instantiate a new object of the aspect class passed as parameter', () => {
         let executeHasBeenCalled = false;
-        class SomeAspect implements Aspect {
-            execute(ctx: AspectContext) {
+        class SomeAspect implements Aspect<void, []> {
+            execute(ctx: AspectContext<void, []>) {
                 executeHasBeenCalled = true;
             }
         }
@@ -79,13 +88,13 @@ describe('UseAspect', () => {
 
     it('should not execute the aspect annotated with Advice.After if an error is thrown in method', () => {
         expect(() => sample.throwError()).toThrow(Error);
-        expect(beforeAspect.execute).toHaveBeenCalledTimes(1);
-        expect(afterAspect.execute).toHaveBeenCalledTimes(0);
+        expect(throwBeforeAspect.execute).toHaveBeenCalledTimes(1);
+        expect(throwAfterAspect.execute).toHaveBeenCalledTimes(0);
     });
 
     describe('for async functions', () => {
-        const aspect = mock<Aspect>();
-        aspect.execute.mockImplementation((ctx: AspectContext) => {
+        const aspect = mock<Aspect<number, []>>();
+        aspect.execute.mockImplementation((ctx: AspectContext<number, []>) => {
             return 42;
         });
 

test/resetAllAspects.test.ts

@@ -8,7 +8,7 @@ import { CalculatorCls } from './samples/CalculatorCls.sample';
 
 describe('resetAllAspects', () => {
     let calculator: CalculatorCls;
-    const aspect = mock<Aspect>();
+    const aspect = mock<Aspect<number, [number, number]>>();
 
     beforeEach(() => {
         jest.clearAllMocks();