feat: infer definition type from "types" argument in getDefOrThrow (#1812)

feat(lex-cli): infer definition type from "types" argument in getDefOrThrow

packages/lex-cli/src/codegen/client.ts

@@ -549,9 +549,7 @@ function genClientXrpcCommon(
   lexicons: Lexicons,
   lexUri: string,
 ) {
-  const def = lexicons.getDefOrThrow(lexUri, ['query', 'procedure']) as
-    | LexXrpcQuery
-    | LexXrpcProcedure
+  const def = lexicons.getDefOrThrow(lexUri, ['query', 'procedure'])
 
   //= export interface CallOptions {...}
   const opts = file.addInterface({
@@ -635,7 +633,7 @@ function genClientRecord(
   lexicons: Lexicons,
   lexUri: string,
 ) {
-  const def = lexicons.getDefOrThrow(lexUri, ['record']) as LexRecord
+  const def = lexicons.getDefOrThrow(lexUri, ['record'])
 
   //= export interface Record {...}
   genObject(file, imports, lexUri, def.record, 'Record')

packages/lex-cli/src/codegen/lex-gen.ts

@@ -7,10 +7,7 @@ import {
   LexArray,
   LexPrimitive,
   LexBlob,
-  LexXrpcProcedure,
-  LexXrpcQuery,
   LexToken,
-  LexXrpcSubscription,
   LexCidLink,
   LexBytes,
   LexIpldType,
@@ -270,7 +267,7 @@ export function genXrpcParams(
     'query',
     'subscription',
     'procedure',
-  ]) as LexXrpcQuery
+  ])
 
   //= export interface QueryParams {...}
   const iface = file.addInterface({
@@ -306,12 +303,9 @@ export function genXrpcInput(
   lexUri: string,
   defaultsArePresent = true,
 ) {
-  const def = lexicons.getDefOrThrow(lexUri, [
-    'query',
-    'procedure',
-  ]) as LexXrpcProcedure
+  const def = lexicons.getDefOrThrow(lexUri, ['query', 'procedure'])
 
-  if (def.input?.schema) {
+  if (def.type === 'procedure' && def.input?.schema) {
     if (def.input.schema.type === 'ref' || def.input.schema.type === 'union') {
       //= export type InputSchema = ...
       const refs =
@@ -340,7 +334,7 @@ export function genXrpcInput(
         defaultsArePresent,
       )
     }
-  } else if (def.input?.encoding) {
+  } else if (def.type === 'procedure' && def.input?.encoding) {
     //= export type InputSchema = string | Uint8Array
     file.addTypeAlias({
       isExported: true,
@@ -368,7 +362,7 @@ export function genXrpcOutput(
     'query',
     'subscription',
     'procedure',
-  ]) as LexXrpcQuery | LexXrpcSubscription | LexXrpcProcedure
+  ])
 
   const schema =
     def.type === 'subscription' ? def.message?.schema : def.output?.schema

packages/lex-cli/src/codegen/server.ts

@@ -4,14 +4,7 @@ import {
   SourceFile,
   VariableDeclarationKind,
 } from 'ts-morph'
-import {
-  Lexicons,
-  LexiconDoc,
-  LexXrpcProcedure,
-  LexXrpcQuery,
-  LexRecord,
-  LexXrpcSubscription,
-} from '@atproto/lexicon'
+import { Lexicons, LexiconDoc } from '@atproto/lexicon'
 import { NSID } from '@atproto/syntax'
 import { gen, lexiconsTs, utilTs } from './common'
 import { GeneratedAPI } from '../types'
@@ -411,9 +404,8 @@ function genServerXrpcMethod(
   lexicons: Lexicons,
   lexUri: string,
 ) {
-  const def = lexicons.getDefOrThrow(lexUri, ['query', 'procedure']) as
-    | LexXrpcQuery
-    | LexXrpcProcedure
+  const def = lexicons.getDefOrThrow(lexUri, ['query', 'procedure'])
+
   file.addImportDeclaration({
     moduleSpecifier: '@atproto/xrpc-server',
     namedImports: [{ name: 'HandlerAuth' }],
@@ -543,9 +535,7 @@ function genServerXrpcStreaming(
   lexicons: Lexicons,
   lexUri: string,
 ) {
-  const def = lexicons.getDefOrThrow(lexUri, [
-    'subscription',
-  ]) as LexXrpcSubscription
+  const def = lexicons.getDefOrThrow(lexUri, ['subscription'])
 
   file.addImportDeclaration({
     moduleSpecifier: '@atproto/xrpc-server',
@@ -601,7 +591,7 @@ function genServerRecord(
   lexicons: Lexicons,
   lexUri: string,
 ) {
-  const def = lexicons.getDefOrThrow(lexUri, ['record']) as LexRecord
+  const def = lexicons.getDefOrThrow(lexUri, ['record'])
 
   //= export interface Record {...}
   genObject(file, imports, lexUri, def.record, 'Record')

packages/lexicon/src/lexicons.ts

@@ -1,8 +1,6 @@
 import {
   LexiconDoc,
   LexRecord,
-  LexXrpcProcedure,
-  LexXrpcQuery,
   LexUserType,
   LexiconDefNotFoundError,
   InvalidLexiconError,
@@ -10,7 +8,6 @@ import {
   ValidationError,
   isObj,
   hasProp,
-  LexXrpcSubscription,
 } from './types'
 import {
   assertValidRecord,
@@ -91,7 +88,14 @@ export class Lexicons {
   /**
    * Get a def, throw if not found. Throws on not found.
    */
-  getDefOrThrow(uri: string, types?: string[]): LexUserType {
+  getDefOrThrow<T extends LexUserType['type'] = LexUserType['type']>(
+    uri: string,
+    types?: readonly T[],
+  ): Extract<LexUserType, { type: T }>
+  getDefOrThrow(
+    uri: string,
+    types?: readonly LexUserType['type'][],
+  ): LexUserType {
     const def = this.getDef(uri)
     if (!def) {
       throw new LexiconDefNotFoundError(`Lexicon not found: ${uri}`)
@@ -154,11 +158,7 @@ export class Lexicons {
       'procedure',
       'subscription',
     ])
-    return assertValidXrpcParams(
-      this,
-      def as LexXrpcProcedure | LexXrpcQuery | LexXrpcSubscription,
-      value,
-    )
+    return assertValidXrpcParams(this, def, value)
   }
 
   /**
@@ -167,7 +167,7 @@ export class Lexicons {
   assertValidXrpcInput(lexUri: string, value: unknown) {
     lexUri = toLexUri(lexUri)
     const def = this.getDefOrThrow(lexUri, ['procedure'])
-    return assertValidXrpcInput(this, def as LexXrpcProcedure, value)
+    return assertValidXrpcInput(this, def, value)
   }
 
   /**
@@ -176,11 +176,7 @@ export class Lexicons {
   assertValidXrpcOutput(lexUri: string, value: unknown) {
     lexUri = toLexUri(lexUri)
     const def = this.getDefOrThrow(lexUri, ['query', 'procedure'])
-    return assertValidXrpcOutput(
-      this,
-      def as LexXrpcProcedure | LexXrpcQuery,
-      value,
-    )
+    return assertValidXrpcOutput(this, def, value)
   }
 
   /**
@@ -189,7 +185,7 @@ export class Lexicons {
   assertValidXrpcMessage<T = unknown>(lexUri: string, value: unknown): T {
     lexUri = toLexUri(lexUri)
     const def = this.getDefOrThrow(lexUri, ['subscription'])
-    return assertValidXrpcMessage(this, def as LexXrpcSubscription, value) as T
+    return assertValidXrpcMessage(this, def, value) as T
   }
 
   /**