From 0f98a43ea6bf3885fec8051d96b0a4b99b8212c4 Mon Sep 17 00:00:00 2001
From: Bhushan Srivastava <59949692+he11owthere@users.noreply.github.com>
Date: Wed, 29 Nov 2023 20:28:31 +0530
Subject: [PATCH] Update docs_formatter.py

---
 ivy_lint/formatters/docs_formatter.py | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/ivy_lint/formatters/docs_formatter.py b/ivy_lint/formatters/docs_formatter.py
index f5c0fe6..f79cf83 100644
--- a/ivy_lint/formatters/docs_formatter.py
+++ b/ivy_lint/formatters/docs_formatter.py
@@ -20,26 +20,29 @@ def format_docstring(self, doc):
         codeblock_start_lines = set()  # This will store indices of lines which start a code block
         lines_to_modify = set()  # This will store the indices of indented lines not containing "..."
         is_codeblock_cont = False
+        lb = 0
+        rb = 0
         
         for idx, line in enumerate(lines):
             stripped_line = line.strip()
-
-            if is_codeblock and is_codeblock_cont and not stripped_line.startswith(('>>>', '...')):
-                lines_to_modify.add(idx)
-            if stripped_line.startswith('>>>'):
-                if "(" in line:
-                    is_codeblock_cont = True
-                #if stripped_line.endswith(')'):
-                if ")" in line:
-                    is_codeblock_cont = False
             if not is_codeblock and stripped_line.startswith('>>>'):
                 is_codeblock = True
                 codeblock_start_lines.add(idx)
-            #elif is_codeblock and (not stripped_line or (not stripped_line.startswith(('>>>', '...', '[', '(')) and not stripped_line.endswith((')', ',', '\\', '(', '[', ']')))):
             elif is_codeblock and not is_codeblock_cont and (not stripped_line or (not stripped_line.startswith(('>>>', '...')))):
                 is_codeblock = False
-            #if is_codeblock and (not stripped_line.startswith(('>>>', '...')) and (stripped_line.startswith(('[', '(')) or stripped_line.endswith((')', ',', '\\', '(', '[', ']')))):
-             #   lines_to_modify.add(idx)
+            if is_codeblock:
+                lb += line.count('(')
+                rb += line.count(')')
+                if rb >= lb:
+                    rb = 0
+                    lb = 0
+                    is_codeblock_cont = False
+                else:
+                    lb = lb - rb
+                    rb = 0
+                    is_codeblock_cont = True
+                    if not stripped_line.startswith(('>>>', '...')):
+                        lines_to_modify.add(idx) 
         
         # Add blank lines before code blocks
         formatted_lines = []