From a0a60458ee806474fd2e6b66cf0160763b9f136c Mon Sep 17 00:00:00 2001
From: Lubos <lubos@heyapi.dev>
Date: Mon, 16 Mar 2026 22:24:11 +0100
Subject: [PATCH] refactor: clean up returns mixin

---
 .../src/py-dsl/mixins/returns.ts               | 18 +++---------------
 1 file changed, 3 insertions(+), 15 deletions(-)

diff --git a/packages/openapi-python/src/py-dsl/mixins/returns.ts b/packages/openapi-python/src/py-dsl/mixins/returns.ts
index 91c976984..896e16ae8 100644
--- a/packages/openapi-python/src/py-dsl/mixins/returns.ts
+++ b/packages/openapi-python/src/py-dsl/mixins/returns.ts
@@ -1,8 +1,6 @@
 import type { AnalysisContext, Node, NodeName } from '@hey-api/codegen-core';
-import { isSymbol } from '@hey-api/codegen-core';
 
 import type { py } from '../../ts-python';
-import { IdPyDsl } from '../expr/identifier';
 import type { BaseCtor, MixinCtor } from './types';
 
 export interface ReturnsMethods extends Node {
@@ -12,7 +10,7 @@ export interface ReturnsMethods extends Node {
 
 export function ReturnsMixin<T extends py.Node, TBase extends BaseCtor<T>>(Base: TBase) {
   abstract class Returns extends Base {
-    protected _returns?: IdPyDsl | py.Expression;
+    protected _returns?: NodeName | py.Expression;
 
     override analyze(ctx: AnalysisContext): void {
       super.analyze(ctx);
@@ -20,22 +18,12 @@ export function ReturnsMixin<T extends py.Node, TBase extends BaseCtor<T>>(Base:
     }
 
     returns(type: NodeName | py.Expression): this {
-      if (typeof type === 'string' || isSymbol(type)) {
-        this._returns = new IdPyDsl(type);
-      } else {
-        this._returns = type as py.Expression;
-      }
+      this._returns = type;
       return this;
     }
 
     protected $returns(): py.Expression | undefined {
-      if (!this._returns) {
-        return;
-      }
-      if (this._returns instanceof IdPyDsl) {
-        return this._returns.toAst();
-      }
-      return this._returns as py.Expression;
+      return this.$node(this._returns);
     }
   }
 
-- 
2.51.2