Remove extra call if it's call expression

Fatal1ty · Nov 15, 2023 · 8597bd9 · 8597bd9
1 parent 66265ab
commit 8597bd9
Showing 1 changed file with 15 additions and 9 deletions.
diff --git a/mashumaro/codecs/_builder.py b/mashumaro/codecs/_builder.py
@@ -1,4 +1,4 @@
-from dataclasses import is_dataclass
+import re
 from typing import Any, Callable, Optional, Type
 
 from mashumaro.core.meta.code.builder import CodeBuilder
@@ -11,6 +11,8 @@
 from mashumaro.core.meta.types.pack import PackerRegistry
 from mashumaro.core.meta.types.unpack import UnpackerRegistry
 
+CALL_EXPR = re.compile(r"^([^ ]+)\(value\)$")
+
 
 class CodecCodeBuilder(CodeBuilder):
     @classmethod
@@ -48,10 +50,12 @@ def add_decode_method(
             )
             self.add_line(f"return {unpacked_value}")
         self.add_line("setattr(decoder_obj, 'decode', decode)")
-        if pre_decoder_func is None and is_dataclass(shape_type):
-            method_name = unpacked_value.partition("(")[0]
-            self.lines.reset()
-            self.add_line(f"setattr(decoder_obj, 'decode', {method_name})")
+        if pre_decoder_func is None:
+            m = CALL_EXPR.match(unpacked_value)
+            if m:
+                method_name = m.group(1)
+                self.lines.reset()
+                self.add_line(f"setattr(decoder_obj, 'decode', {method_name})")
         self.ensure_object_imported(decoder_obj, "decoder_obj")
         self.ensure_object_imported(self.cls, "cls")
         self.compile()
@@ -89,10 +93,12 @@ def add_encode_method(
             else:
                 self.add_line(f"return {packed_value}")
         self.add_line("setattr(encoder_obj, 'encode', encode)")
-        if post_encoder_func is None and is_dataclass(shape_type):
-            method_name = packed_value.partition("(")[0]
-            self.lines.reset()
-            self.add_line(f"setattr(encoder_obj, 'encode', {method_name})")
+        if post_encoder_func is None:
+            m = CALL_EXPR.match(packed_value)
+            if m:
+                method_name = m.group(1)
+                self.lines.reset()
+                self.add_line(f"setattr(encoder_obj, 'encode', {method_name})")
         self.ensure_object_imported(encoder_obj, "encoder_obj")
         self.ensure_object_imported(self.cls, "cls")
         self.ensure_object_imported(self.cls, "self")