Merge pull request rocky#502 from gdesmar/docstring_bytes

rocky · web-flow · commit 27c869b69a09 · 2024-10-16T20:31:12.000-04:00
Fix for print_docstring()'s `docstring.find(quote)` Type error
diff --git a/test/bytecode_2.7/16_bytestring_docstring.pyc b/test/bytecode_2.7/16_bytestring_docstring.pyc
diff --git a/test/bytecode_3.8/16_no_bytestring_docstring.pyc b/test/bytecode_3.8/16_no_bytestring_docstring.pyc
diff --git a/test/simple_source/stmts/16_bytestring_docstring.py b/test/simple_source/stmts/16_bytestring_docstring.py
@@ -0,0 +1,45 @@
+"""Module docstring"""
+class A:
+    b"""Got \xe7\xfe Bytes?"""
+    assert __doc__ == b"""Got \xe7\xfe Bytes?"""
+
+    def class_func(self):
+       b"""Got \xe7\xfe Bytes?"""
+       assert __doc__ == """Module docstring"""
+
+class B:
+    """Got no Bytes?"""
+    assert __doc__ == """Got no Bytes?"""
+
+    def class_func(self):
+        """Got no Bytes?"""
+        assert __doc__ == """Module docstring"""
+
+def single_func():
+    """single docstring?"""
+    assert __doc__ == """Module docstring"""
+
+def single_byte_func():
+    b"""Got \xe7\xfe Bytes?"""
+    assert __doc__ == """Module docstring"""
+
+assert __doc__ == """Module docstring"""
+
+assert single_func.__doc__ == """single docstring?"""
+single_func()
+
+assert single_byte_func.__doc__ == b"""Got \xe7\xfe Bytes?"""
+single_byte_func()
+
+assert A.__doc__ == b"""Got \xe7\xfe Bytes?"""
+assert A.class_func.__doc__ == b"""Got \xe7\xfe Bytes?"""
+a = A()
+assert a.class_func.__doc__ == b"""Got \xe7\xfe Bytes?"""
+a.class_func()
+
+assert B.__doc__ == """Got no Bytes?"""
+assert B.class_func.__doc__ == """Got no Bytes?"""
+b = B()
+assert b.class_func.__doc__ == """Got no Bytes?"""
+b.class_func()
+
diff --git a/test/simple_source/stmts/16_no_bytestring_docstring.py b/test/simple_source/stmts/16_no_bytestring_docstring.py
@@ -0,0 +1,45 @@
+"""Module docstring"""
+class A:
+    b"""Got \xe7\xfe Bytes?"""
+    assert __doc__ == """Module docstring"""
+
+    def class_func(self):
+       b"""Got \xe7\xfe Bytes?"""
+       assert __doc__ == """Module docstring"""
+
+class B:
+    """Got no Bytes?"""
+    assert __doc__ == """Got no Bytes?"""
+
+    def class_func(self):
+        """Got no Bytes?"""
+        assert __doc__ == """Module docstring"""
+
+def single_func():
+    """single docstring?"""
+    assert __doc__ == """Module docstring"""
+
+def single_byte_func():
+    b"""Got \xe7\xfe Bytes?"""
+    assert __doc__ == """Module docstring"""
+
+assert __doc__ == """Module docstring"""
+
+assert single_func.__doc__ == """single docstring?"""
+single_func()
+
+assert single_byte_func.__doc__ is None
+single_byte_func()
+
+assert A.__doc__ is None
+assert A.class_func.__doc__ is None
+a = A()
+assert a.class_func.__doc__ is None
+a.class_func()
+
+assert B.__doc__ == """Got no Bytes?"""
+assert B.class_func.__doc__ == """Got no Bytes?"""
+b = B()
+assert b.class_func.__doc__ == """Got no Bytes?"""
+b.class_func()
+
diff --git a/uncompyle6/semantics/helper.py b/uncompyle6/semantics/helper.py
@@ -152,6 +152,9 @@ def is_lambda_mode(compile_mode: str) -> bool:
 
 
 def print_docstring(self, indent, docstring):
+    if isinstance(docstring, bytes):
+        docstring = docstring.decode("utf8", errors="backslashreplace")
+
     quote = '"""'
     if docstring.find(quote) >= 0:
         if docstring.find("'''") == -1:
diff --git a/uncompyle6/semantics/n_actions.py b/uncompyle6/semantics/n_actions.py
@@ -26,7 +26,7 @@
     PRECEDENCE,
     minint,
 )
-from uncompyle6.semantics.helper import find_code_node, flatten_list
+from uncompyle6.semantics.helper import find_code_node, flatten_list, print_docstring
 from uncompyle6.util import better_repr, get_code_name
 
 
@@ -541,70 +541,7 @@ def n_docstring(self, node):
         else:
             docstring = node[0].pattr
 
-        quote = '"""'
-        if docstring.find(quote) >= 0:
-            if docstring.find("'''") == -1:
-                quote = "'''"
-
-        self.write(indent)
-        docstring = repr(docstring.expandtabs())[1:-1]
-
-        for orig, replace in (
-            ("\\\\", "\t"),
-            ("\\r\\n", "\n"),
-            ("\\n", "\n"),
-            ("\\r", "\n"),
-            ('\\"', '"'),
-            ("\\'", "'"),
-        ):
-            docstring = docstring.replace(orig, replace)
-
-        # Do a raw string if there are backslashes but no other escaped characters:
-        # also check some edge cases
-        if (
-            "\t" in docstring
-            and "\\" not in docstring
-            and len(docstring) >= 2
-            and docstring[-1] != "\t"
-            and (docstring[-1] != '"' or docstring[-2] == "\t")
-        ):
-            self.write("r")  # raw string
-            # Restore backslashes unescaped since raw
-            docstring = docstring.replace("\t", "\\")
-        else:
-            # Escape the last character if it is the same as the
-            # triple quote character.
-            quote1 = quote[-1]
-            if len(docstring) and docstring[-1] == quote1:
-                docstring = docstring[:-1] + "\\" + quote1
-
-            # Escape triple quote when needed
-            if quote == '"""':
-                replace_str = '\\"""'
-            else:
-                assert quote == "'''"
-                replace_str = "\\'''"
-
-            docstring = docstring.replace(quote, replace_str)
-            docstring = docstring.replace("\t", "\\\\")
-
-        lines = docstring.split("\n")
-
-        self.write(quote)
-        if len(lines) == 0:
-            self.println(quote)
-        elif len(lines) == 1:
-            self.println(lines[0], quote)
-        else:
-            self.println(lines[0])
-            for line in lines[1:-1]:
-                if line:
-                    self.println(line)
-                else:
-                    self.println("\n\n")
-                    pass
-                pass
-            self.println(lines[-1], quote)
+        print_docstring(self, indent, docstring)
         self.prune()
 
     def n_elifelsestmtr(self, node: SyntaxTree):
