add register codec functions

openrewrite · Oct 31, 2024 · 42c60a3 · 42c60a3
1 parent d4da5b2
commit 42c60a3
Show file tree

Hide file tree

Showing 8 changed files with 208 additions and 116 deletions.
diff --git a/openrewrite/src/index.ts b/openrewrite/src/index.ts
@@ -1 +1 @@
-export {}
+export * from './types/register-codex';
diff --git a/openrewrite/src/java/remote/index.ts b/openrewrite/src/java/remote/index.ts
@@ -1,5 +1,3 @@
-// import for side-effects
-import './register';
-
+export * from './register';
 export * from './receiver';
 export * from './sender';
diff --git a/openrewrite/src/java/remote/register.ts b/openrewrite/src/java/remote/register.ts
@@ -1,129 +1,205 @@
-import {RemotingContext} from "@openrewrite/rewrite-remote";
-import {isJava, JavaType} from "../tree";
-import {DeserializationContext, ReceiverContext} from "@openrewrite/rewrite-remote";
-import {CborDecoder, CborEncoder} from "@jsonjoy.com/json-pack/lib/cbor";
-import {ValueType} from "@openrewrite/rewrite-remote";
-import {SenderContext, SerializationContext} from "@openrewrite/rewrite-remote";
-import {JavaSender} from "./sender";
-import {JavaReceiver} from "./receiver";
+import { isJava, JavaType } from '../tree';
+import { CborDecoder, CborEncoder } from '@jsonjoy.com/json-pack/lib/cbor';
+import { ValueType } from '@openrewrite/rewrite-remote';
+import { SerializationContext } from '@openrewrite/rewrite-remote';
+import { JavaSender } from './sender';
+import { JavaReceiver } from './receiver';
 
-console.log("registering java codecs");
+import type { DeserializationContext } from '@openrewrite/rewrite-remote';
+import type { RegisterCodecs } from '../../types/register-codex';
 
-SenderContext.register(isJava, () => new JavaSender());
-ReceiverContext.register(isJava, () => new JavaReceiver());
+export const registerCodecs: RegisterCodecs = (
+  senderContext,
+  receiverContext,
+  remotingContext
+) => {
+  console.log('registering java codecs');
 
-RemotingContext.registerValueDeserializer(JavaType.Class, function (type: any, decoder: CborDecoder, context: DeserializationContext): JavaType.Class {
-    let cls: {
-        [key: string]: any
-    } = type == JavaType.ShallowClass ? new JavaType.ShallowClass() : new JavaType.Class();
-    while (decoder.reader.peak() !== 255) {
+  senderContext?.register(isJava, () => new JavaSender());
+  receiverContext?.register(isJava, () => new JavaReceiver());
+
+  remotingContext?.registerValueDeserializer(
+    JavaType.Class,
+    function (
+      type: any,
+      decoder: CborDecoder,
+      context: DeserializationContext
+    ): JavaType.Class {
+      let cls: {
+        [key: string]: any;
+      } =
+        type == JavaType.ShallowClass
+          ? new JavaType.ShallowClass()
+          : new JavaType.Class();
+      while (decoder.reader.peak() !== 255) {
         const key = decoder.key();
         if (key === '@ref') {
-            context.remotingContext.addById(decoder.val() as number, cls);
+          context.remotingContext.addById(decoder.val() as number, cls);
         } else {
-            cls['_' + key] = decoder.val();
+          cls['_' + key] = decoder.val();
         }
+      }
+      decoder.reader.x++;
+      return cls as JavaType.Class;
     }
-    decoder.reader.x++;
-    return cls as JavaType.Class;
-});
+  );
 
-RemotingContext.registerValueDeserializer(JavaType.Primitive, function (type: any, decoder: CborDecoder, context: DeserializationContext): JavaType.Primitive {
-    const key = decoder.key();
-    let kind: JavaType.PrimitiveKind;
-    if (key === '_kind') {
+  remotingContext?.registerValueDeserializer(
+    JavaType.Primitive,
+    function (
+      type: any,
+      decoder: CborDecoder,
+      context: DeserializationContext
+    ): JavaType.Primitive {
+      const key = decoder.key();
+      let kind: JavaType.PrimitiveKind;
+      if (key === '_kind') {
         kind = decoder.val() as JavaType.PrimitiveKind;
-    } else {
+      } else {
         throw new Error('Unexpected key: ' + key);
+      }
+      decoder.reader.x++;
+      return JavaType.Primitive.of(kind);
     }
-    decoder.reader.x++;
-    return JavaType.Primitive.of(kind);
-});
+  );
 
-RemotingContext.registerValueSerializer(JavaType.Primitive, function (value: JavaType.Primitive, type: ValueType, typeName: string | null, encoder: CborEncoder, context: SerializationContext): void {
-    encoder.writeArrHdr(2);
-    encoder.writeStr("org.openrewrite.java.tree.JavaType$Primitive");
-    encoder.writeUInteger(value.kind);
-});
+  remotingContext?.registerValueSerializer(
+    JavaType.Primitive,
+    function (
+      value: JavaType.Primitive,
+      type: ValueType,
+      typeName: string | null,
+      encoder: CborEncoder,
+      context: SerializationContext
+    ): void {
+      encoder.writeArrHdr(2);
+      encoder.writeStr('org.openrewrite.java.tree.JavaType$Primitive');
+      encoder.writeUInteger(value.kind);
+    }
+  );
 
-RemotingContext.registerValueSerializer(JavaType.Class, function (value: JavaType.Class, type: ValueType, typeName: string | null, encoder: CborEncoder, context: SerializationContext): void {
-    const id = context.remotingContext.tryGetId(value);
-    if (id !== undefined) {
+  remotingContext?.registerValueSerializer(
+    JavaType.Class,
+    function (
+      value: JavaType.Class,
+      type: ValueType,
+      typeName: string | null,
+      encoder: CborEncoder,
+      context: SerializationContext
+    ): void {
+      const id = context.remotingContext.tryGetId(value);
+      if (id !== undefined) {
         encoder.writeUInteger(id);
         return;
-    }
-    encoder.writeStartObj();
-    encoder.writeStr('@c');
-    encoder.writeStr(value instanceof JavaType.ShallowClass ? "org.openrewrite.java.tree.JavaType$ShallowClass" : "org.openrewrite.java.tree.JavaType$Class");
-    encoder.writeStr('@ref');
-    encoder.writeUInteger(context.remotingContext.add(value));
-    encoder.writeStr('flagsBitMap');
-    encoder.writeNumber(value.flagsBitMap);
-    encoder.writeStr('fullyQualifiedName');
-    encoder.writeStr(value.fullyQualifiedName);
-    encoder.writeStr('kind');
-    encoder.writeUInteger(value.kind);
-    if (value.typeParameters) {
+      }
+      encoder.writeStartObj();
+      encoder.writeStr('@c');
+      encoder.writeStr(
+        value instanceof JavaType.ShallowClass
+          ? 'org.openrewrite.java.tree.JavaType$ShallowClass'
+          : 'org.openrewrite.java.tree.JavaType$Class'
+      );
+      encoder.writeStr('@ref');
+      encoder.writeUInteger(context.remotingContext.add(value));
+      encoder.writeStr('flagsBitMap');
+      encoder.writeNumber(value.flagsBitMap);
+      encoder.writeStr('fullyQualifiedName');
+      encoder.writeStr(value.fullyQualifiedName);
+      encoder.writeStr('kind');
+      encoder.writeUInteger(value.kind);
+      if (value.typeParameters) {
         encoder.writeStr('typeParameters');
         encoder.writeArrHdr(value.typeParameters.length);
-        value.typeParameters.forEach(t => context.serialize(t, ValueType.Object, null, encoder));
-    }
-    if (value.supertype) {
+        value.typeParameters.forEach((t) =>
+          context.serialize(t, ValueType.Object, null, encoder)
+        );
+      }
+      if (value.supertype) {
         encoder.writeStr('supertype');
         context.serialize(value.supertype, ValueType.Object, null, encoder);
-    }
-    if (value.owningClass) {
+      }
+      if (value.owningClass) {
         encoder.writeStr('owningClass');
         context.serialize(value.owningClass, ValueType.Object, null, encoder);
-    }
-    if (value.annotations) {
+      }
+      if (value.annotations) {
         encoder.writeStr('annotations');
         encoder.writeArrHdr(value.annotations.length);
-        value.annotations.forEach(t => context.serialize(t, ValueType.Object, null, encoder));
-    }
-    if (value.interfaces) {
+        value.annotations.forEach((t) =>
+          context.serialize(t, ValueType.Object, null, encoder)
+        );
+      }
+      if (value.interfaces) {
         encoder.writeStr('interfaces');
         encoder.writeArrHdr(value.interfaces.length);
-        value.interfaces.forEach(t => context.serialize(t, ValueType.Object, null, encoder));
-    }
-    if (value.members) {
+        value.interfaces.forEach((t) =>
+          context.serialize(t, ValueType.Object, null, encoder)
+        );
+      }
+      if (value.members) {
         encoder.writeStr('members');
         encoder.writeArrHdr(value.members.length);
-        value.members.forEach(t => context.serialize(t, ValueType.Object, null, encoder));
-    }
-    if (value.methods) {
+        value.members.forEach((t) =>
+          context.serialize(t, ValueType.Object, null, encoder)
+        );
+      }
+      if (value.methods) {
         encoder.writeStr('methods');
         encoder.writeArrHdr(value.methods.length);
-        value.methods.forEach(t => context.serialize(t, ValueType.Object, null, encoder));
+        value.methods.forEach((t) =>
+          context.serialize(t, ValueType.Object, null, encoder)
+        );
+      }
+      encoder.writeEndObj();
     }
-    encoder.writeEndObj();
-});
+  );
 
-RemotingContext.registerValueSerializer(JavaType.Union, function (value: JavaType.Union, type: ValueType, typeName: string | null, encoder: CborEncoder, context: SerializationContext): void {
-    const id = context.remotingContext.tryGetId(value);
-    if (id !== undefined) {
+  remotingContext?.registerValueSerializer(
+    JavaType.Union,
+    function (
+      value: JavaType.Union,
+      type: ValueType,
+      typeName: string | null,
+      encoder: CborEncoder,
+      context: SerializationContext
+    ): void {
+      const id = context.remotingContext.tryGetId(value);
+      if (id !== undefined) {
         encoder.writeUInteger(id);
         return;
+      }
+      encoder.writeMapHdr(3);
+      encoder.writeStr('@c');
+      encoder.writeStr('org.openrewrite.java.tree.JavaType$MultiCatch');
+      encoder.writeStr('@ref');
+      encoder.writeUInteger(context.remotingContext.add(value));
+      encoder.writeStr('throwableTypes');
+      encoder.writeArrHdr(value.types.length);
+      value.types.forEach((t) =>
+        context.serialize(t, ValueType.Object, null, encoder)
+      );
     }
-    encoder.writeMapHdr(3);
-    encoder.writeStr('@c');
-    encoder.writeStr("org.openrewrite.java.tree.JavaType$MultiCatch");
-    encoder.writeStr('@ref');
-    encoder.writeUInteger(context.remotingContext.add(value));
-    encoder.writeStr('throwableTypes');
-    encoder.writeArrHdr(value.types.length);
-    value.types.forEach(t => context.serialize(t, ValueType.Object, null, encoder));
-});
+  );
 
-RemotingContext.registerValueSerializer(JavaType.Unknown, function (value: JavaType.Unknown, type: ValueType, typeName: string | null, encoder: CborEncoder, context: SerializationContext): void {
-    const id = context.remotingContext.tryGetId(value);
-    if (id !== undefined) {
+  remotingContext?.registerValueSerializer(
+    JavaType.Unknown,
+    function (
+      value: JavaType.Unknown,
+      type: ValueType,
+      typeName: string | null,
+      encoder: CborEncoder,
+      context: SerializationContext
+    ): void {
+      const id = context.remotingContext.tryGetId(value);
+      if (id !== undefined) {
         encoder.writeUInteger(id);
         return;
+      }
+      encoder.writeMapHdr(2);
+      encoder.writeStr('@c');
+      encoder.writeStr('org.openrewrite.java.tree.JavaType$Unknown');
+      encoder.writeStr('@ref');
+      encoder.writeUInteger(context.remotingContext.add(value));
     }
-    encoder.writeMapHdr(2);
-    encoder.writeStr('@c');
-    encoder.writeStr("org.openrewrite.java.tree.JavaType$Unknown");
-    encoder.writeStr('@ref');
-    encoder.writeUInteger(context.remotingContext.add(value));
-});
+  );
+};
diff --git a/openrewrite/src/javascript/remote/index.ts b/openrewrite/src/javascript/remote/index.ts
@@ -1,5 +1,3 @@
-// import for side-effects
-import './register';
-
+export * from './register';
 export * from './receiver';
 export * from './sender';
diff --git a/openrewrite/src/javascript/remote/register.ts b/openrewrite/src/javascript/remote/register.ts
@@ -1,8 +1,14 @@
-import {isJavaScript} from "../tree";
-import {ReceiverContext} from "@openrewrite/rewrite-remote";
-import {SenderContext} from "@openrewrite/rewrite-remote";
-import {JavaScriptSender} from "./sender";
-import {JavaScriptReceiver} from "./receiver";
-console.log("registering javascript codecs");
-SenderContext.register(isJavaScript, () => new JavaScriptSender());
-ReceiverContext.register(isJavaScript, () => new JavaScriptReceiver());
+import { isJavaScript } from '../tree';
+import { JavaScriptSender } from './sender';
+import { JavaScriptReceiver } from './receiver';
+import type { RegisterCodecs } from '../../types/register-codex';
+
+export const registerCodecs: RegisterCodecs = (
+  senderContext,
+  receiverContext,
+  remotingContext
+) => {
+  console.log('registering javascript codecs');
+  senderContext?.register(isJavaScript, () => new JavaScriptSender());
+  receiverContext?.register(isJavaScript, () => new JavaScriptReceiver());
+};
diff --git a/openrewrite/src/types/register-codex.ts b/openrewrite/src/types/register-codex.ts
@@ -0,0 +1,11 @@
+import type {
+  ReceiverContext,
+  RemotingContext,
+  SenderContext,
+} from '@openrewrite/rewrite-remote';
+
+export type RegisterCodecs = (
+  senderContext?: typeof SenderContext,
+  receiverContext?: typeof ReceiverContext,
+  remotingContext?: typeof RemotingContext
+) => void;
diff --git a/openrewrite/src/yaml/remote/index.ts b/openrewrite/src/yaml/remote/index.ts
@@ -1,5 +1,3 @@
-// import for side-effects
-import './register';
-
+export * from './register';
 export * from './receiver';
 export * from './sender';
diff --git a/openrewrite/src/yaml/remote/register.ts b/openrewrite/src/yaml/remote/register.ts
@@ -1,9 +1,14 @@
-import {isYaml} from "../tree";
-import {ReceiverContext} from "@openrewrite/rewrite-remote";
-import {SenderContext} from "@openrewrite/rewrite-remote";
-import {YamlSender} from "./sender";
-import {YamlReceiver} from "./receiver";
+import { isYaml } from '../tree';
+import { YamlSender } from './sender';
+import { YamlReceiver } from './receiver';
+import type { RegisterCodecs } from '../../types/register-codex';
 
-console.log("registering yaml codecs");
-SenderContext.register(isYaml, () => new YamlSender());
-ReceiverContext.register(isYaml, () => new YamlReceiver());
+export const registerCodecs: RegisterCodecs = (
+  senderContext,
+  receiverContext,
+  remotingContext
+) => {
+  console.log('registering yaml codecs');
+  senderContext?.register(isYaml, () => new YamlSender());
+  receiverContext?.register(isYaml, () => new YamlReceiver());
+};