Simple Vim Implementation in Python

vim.py

import curses
import os
import sys

# Modes
NORMAL_MODE = 0
INSERT_MODE = 1
COMMAND_MODE = 2

class PyVim:
    def __init__(self, stdscr, filename=None):
        self.stdscr = stdscr
        self.height, self.width = self.stdscr.getmaxyx()

        self.lines = [""]  # Document content as a list of strings
        self.cursor_y = 0  # Buffer line index
        self.cursor_x = 0  # Buffer column index (character position)
        self.top_row = 0   # First buffer line displayed on screen
        self.left_col = 0  # First buffer column displayed (for potential horizontal scrolling, simplified for now)

        self.mode = NORMAL_MODE
        self.filename = filename
        self.status_message = ""
        self.command_buffer = ""
        self.dirty = False # True if there are unsaved changes

        if self.filename:
            self._load_file(self.filename)
        else:
            self.filename = "[No Name]"
            self.status_message = "New buffer"
            
        self._ensure_cursor_bounds()

    def _load_file(self, filename):
        try:
            if os.path.exists(filename):
                with open(filename, 'r', encoding='utf-8') as f:
                    self.lines = [line.rstrip('\n') for line in f.readlines()]
                if not self.lines: # Ensure at least one empty line if file is empty
                    self.lines = [""]
                self.status_message = f"Loaded \"{filename}\""
            else:
                self.lines = [""]
                self.status_message = f"New file \"{filename}\""
            self.dirty = False
        except Exception as e:
            self.lines = [""] # Fallback to empty buffer
            self.status_message = f"Error loading file: {e}"
            self.dirty = False # Treat as a new buffer if load failed

    def _save_file(self, filename_to_save=None):
        current_filename = filename_to_save if filename_to_save else (self.filename if self.filename != "[No Name]" else None)

        if not current_filename:
            self.status_message = "E32: No file name"
            return False

        try:
            with open(current_filename, 'w', encoding='utf-8') as f:
                for line in self.lines:
                    f.write(line + '\n')
            self.filename = current_filename # Update filename if it was newly provided
            self.status_message = f"\"{current_filename}\" {len(self.lines)}L written"
            self.dirty = False
            return True
        except Exception as e:
            self.status_message = f"Error saving file: {e}"
            return False

    def _ensure_cursor_bounds(self):
        # Ensure cursor_y is within document bounds
        self.cursor_y = max(0, min(self.cursor_y, len(self.lines) - 1))

        # Ensure cursor_x is within line bounds
        current_line_len = len(self.lines[self.cursor_y])
        if self.mode == NORMAL_MODE:
            # In normal mode, cursor_x can be on the last char, or 0 if line is empty
            self.cursor_x = max(0, min(self.cursor_x, max(0, current_line_len -1)))
        else: # INSERT_MODE or COMMAND_MODE (though cursor isn't in text for command)
            # In insert mode, cursor_x can be one past the end of the line
             self.cursor_x = max(0, min(self.cursor_x, current_line_len))

    def _scroll_if_needed(self):
        screen_text_height = self.height - 2 # 1 for status, 1 for command (or future use)

        if self.cursor_y < self.top_row:
            self.top_row = self.cursor_y
        elif self.cursor_y >= self.top_row + screen_text_height:
            self.top_row = self.cursor_y - screen_text_height + 1
        
        # Basic horizontal scroll adjustment (very simplified)
        # This ensures the cursor is visible if it goes off screen to the right
        # A more robust solution would track left_col for viewport
        if self.cursor_x >= self.left_col + self.width -5: # -5 for line numbers and margin
             self.left_col = self.cursor_x - (self.width - 5) + 1
        elif self.cursor_x < self.left_col:
             self.left_col = self.cursor_x
        self.left_col = max(0, self.left_col)


    def _draw_status_bar(self):
        self.stdscr.attron(curses.A_REVERSE)
        status_text = ""
        if self.mode == NORMAL_MODE:
            status_text = "-- NORMAL --"
        elif self.mode == INSERT_MODE:
            status_text = "-- INSERT --"
        elif self.mode == COMMAND_MODE:
            status_text = f":{self.command_buffer}"
        
        file_status = f"{self.filename}{' [+]' if self.dirty else ''}"
        pos_status = f"{self.cursor_y + 1},{self.cursor_x + 1}"
        
        left_part = f"{status_text} {file_status}".ljust((self.width // 3) * 2)
        right_part = pos_status.rjust(self.width - len(left_part))
        
        full_status = (left_part + right_part)[:self.width-1]


        self.stdscr.addstr(self.height - 2, 0, full_status + " " * (self.width - len(full_status) -1))
        self.stdscr.attroff(curses.A_REVERSE)

        # Command/Message line
        msg_line_y = self.height -1
        self.stdscr.addstr(msg_line_y, 0, " " * (self.width -1)) # Clear message line
        if self.mode == COMMAND_MODE:
            self.stdscr.addstr(msg_line_y, 0, f":{self.command_buffer}")
        else:
            self.stdscr.addstr(msg_line_y, 0, self.status_message[:self.width-1])
        
        # Clear status_message after displaying it once, unless it's a persistent one
        if self.status_message and not (self.status_message.startswith("E32:") or self.status_message.startswith("New") or self.status_message.startswith("Loaded")):
             pass # Keep some messages like errors or load status
        # self.status_message = "" # Commented out to keep messages sticky for a bit

    def _draw_text_area(self):
        screen_text_height = self.height - 2 # For status bar and command line
        line_num_width = 4 # Width for line numbers (e.g., "123 ")

        for i in range(screen_text_height):
            buffer_line_idx = self.top_row + i
            screen_y = i

            # Clear the line first
            self.stdscr.addstr(screen_y, 0, " " * (self.width -1))

            if buffer_line_idx < len(self.lines):
                # Display line number
                self.stdscr.attron(curses.color_pair(1)) # Assuming color pair 1 for line numbers
                self.stdscr.addstr(screen_y, 0, str(buffer_line_idx + 1).rjust(line_num_width -1) + " ")
                self.stdscr.attroff(curses.color_pair(1))

                # Display line content, respecting self.left_col for horizontal scrolling
                line_content = self.lines[buffer_line_idx]
                display_string = line_content[self.left_col : self.left_col + self.width - line_num_width -1] # -1 for safety
                self.stdscr.addstr(screen_y, line_num_width, display_string)
            else:
                if buffer_line_idx == 0 and not self.lines[0]: # Special case for truly empty buffer
                     pass # Don't show tilde if it's an empty buffer first line
                elif self.filename != "[No Name]" or len(self.lines) > 1 or self.lines[0]: # Show tilde if not empty or named
                    self.stdscr.attron(curses.color_pair(1))
                    self.stdscr.addstr(screen_y, 0, "~")
                    self.stdscr.attroff(curses.color_pair(1))


    def _render(self):
        self.stdscr.erase() # Clear screen
        self._scroll_if_needed()
        self._draw_text_area()
        self._draw_status_bar()
        
        # Place cursor
        # Screen cursor position needs to be relative to top_row and left_col
        screen_cursor_y = self.cursor_y - self.top_row
        screen_cursor_x = (self.cursor_x - self.left_col) + 4 # +4 for line numbers area

        # Ensure screen cursor is within visible text area bounds before moving
        if 0 <= screen_cursor_y < (self.height - 2) and \
           4 <= screen_cursor_x < self.width: # line_num_width used as offset
            try:
                self.stdscr.move(screen_cursor_y, screen_cursor_x)
            except curses.error:
                # This can happen if cursor is off screen, e.g. during rapid changes.
                # Fallback to a safe position or just ignore.
                # For now, move to a known safe spot if error.
                self.stdscr.move(0,4) # Fallback to top-left of text area

        self.stdscr.refresh()

    def _handle_normal_mode_input(self, key):
        self.status_message = "" # Clear previous messages
        if key == ord('h'):
            if self.cursor_x > 0:
                self.cursor_x -= 1
        elif key == ord('l'):
            current_line_len = len(self.lines[self.cursor_y])
            if current_line_len == 0: # Empty line
                self.cursor_x = 0
            elif self.cursor_x < current_line_len -1 : # -1 because x is 0-indexed for content
                self.cursor_x += 1
        elif key == ord('k'):
            if self.cursor_y > 0:
                self.cursor_y -= 1
                # Try to maintain column, or go to end of shorter line
                current_line_len = len(self.lines[self.cursor_y])
                self.cursor_x = min(self.cursor_x, max(0, current_line_len -1))
                if current_line_len == 0: self.cursor_x = 0

        elif key == ord('j'):
            if self.cursor_y < len(self.lines) - 1:
                self.cursor_y += 1
                # Try to maintain column, or go to end of shorter line
                current_line_len = len(self.lines[self.cursor_y])
                self.cursor_x = min(self.cursor_x, max(0, current_line_len -1))
                if current_line_len == 0: self.cursor_x = 0
        
        elif key == ord('i'): # Insert before cursor
            self.mode = INSERT_MODE
        elif key == ord('a'): # Append after cursor
            self.mode = INSERT_MODE
            current_line_len = len(self.lines[self.cursor_y])
            if current_line_len > 0 : # Only advance if line is not empty
                 self.cursor_x = min(self.cursor_x + 1, current_line_len)
        elif key == ord('o'): # Open line below
            self.mode = INSERT_MODE
            current_line = self.lines[self.cursor_y]
            # For 'o', new line is inserted below, content after cursor on current line stays
            self.lines.insert(self.cursor_y + 1, "")
            self.cursor_y += 1
            self.cursor_x = 0
            self.dirty = True
        elif key == ord('O'): # Open line above
            self.mode = INSERT_MODE
            self.lines.insert(self.cursor_y, "")
            # cursor_y remains the same (it's now the new empty line)
            self.cursor_x = 0
            self.dirty = True

        elif key == ord('x'): # Delete character under cursor
            if self.lines[self.cursor_y]: # If line is not empty
                line = self.lines[self.cursor_y]
                # Ensure cursor_x is valid for deletion
                self.cursor_x = min(self.cursor_x, len(line) -1)
                if self.cursor_x < 0: self.cursor_x = 0 # Should not happen if line not empty

                if len(line) > 0:
                    self.lines[self.cursor_y] = line[:self.cursor_x] + line[self.cursor_x+1:]
                    self.dirty = True
                # Cursor remains, or moves left if it was at the end of the deleted part
                if self.cursor_x >= len(self.lines[self.cursor_y]) and len(self.lines[self.cursor_y]) > 0:
                    self.cursor_x = len(self.lines[self.cursor_y]) -1
                elif not self.lines[self.cursor_y]: # Line became empty
                    self.cursor_x = 0


        elif key == ord('d'): # Potential 'dd'
            self.stdscr.nodelay(True) # Make next getch non-blocking
            next_key = self.stdscr.getch()
            self.stdscr.nodelay(False)
            if next_key == ord('d'): # 'dd' confirmed
                if self.lines: # If there are any lines
                    del self.lines[self.cursor_y]
                    if not self.lines: # If all lines deleted
                        self.lines = [""] # Add back one empty line
                        self.cursor_y = 0
                    else: # Ensure cursor_y is still valid
                        self.cursor_y = min(self.cursor_y, len(self.lines) -1)
                    self.cursor_x = 0 # Move cursor to start of line
                    self.dirty = True
            else: # Not 'dd', might be other 'd' command or nothing
                if next_key != -1: # If a key was pressed, unget it for next loop
                    curses.ungetch(next_key)
                self.status_message = "Unknown 'd' command"

        elif key == ord(':'):
            self.mode = COMMAND_MODE
            self.command_buffer = ""

        self._ensure_cursor_bounds()
        return True # Continue running

    def _handle_insert_mode_input(self, key):
        self.status_message = ""
        current_line = self.lines[self.cursor_y]

        if key == 27: # Escape key
            self.mode = NORMAL_MODE
            # In Vim, Esc moves cursor left if not at BOL in insert mode
            if self.cursor_x > 0:
                self.cursor_x -=1
        elif key == curses.KEY_BACKSPACE or key == 127 or key == 8: # Backspace
            if self.cursor_x > 0:
                self.lines[self.cursor_y] = current_line[:self.cursor_x-1] + current_line[self.cursor_x:]
                self.cursor_x -= 1
                self.dirty = True
            elif self.cursor_y > 0: # Backspace at beginning of line, join with previous
                prev_line_len = len(self.lines[self.cursor_y -1])
                self.lines[self.cursor_y -1] += current_line
                del self.lines[self.cursor_y]
                self.cursor_y -= 1
                self.cursor_x = prev_line_len
                self.dirty = True
        elif key == curses.KEY_ENTER or key == 10 or key == 13: # Enter
            # Split line at cursor
            second_part = current_line[self.cursor_x:]
            self.lines[self.cursor_y] = current_line[:self.cursor_x]
            self.cursor_y += 1
            self.lines.insert(self.cursor_y, second_part)
            self.cursor_x = 0
            self.dirty = True
        elif curses.KEY_MIN <= key <= curses.KEY_MAX: # Arrow keys etc. - ignore in insert for simplicity or handle later
            pass # For now, ignore special keys like arrows in insert mode
        elif 32 <= key <= 126 or key > 127: # Printable characters (basic ASCII and potentially beyond)
            try:
                char_to_insert = chr(key)
                self.lines[self.cursor_y] = current_line[:self.cursor_x] + char_to_insert + current_line[self.cursor_x:]
                self.cursor_x += 1
                self.dirty = True
            except ValueError: # chr(key) might fail for some values from getch if not proper char
                self.status_message = f"Invalid character input: {key}"

        self._ensure_cursor_bounds() # After modification
        return True

    def _handle_command_mode_input(self, key):
        if key == 27: # Escape
            self.mode = NORMAL_MODE
            self.command_buffer = ""
            self.status_message = ""
        elif key == curses.KEY_BACKSPACE or key == 127 or key == 8:
            if self.command_buffer:
                self.command_buffer = self.command_buffer[:-1]
        elif key == curses.KEY_ENTER or key == 10 or key == 13:
            self._execute_command()
            # _execute_command might change mode or request quit
            if self.mode == COMMAND_MODE: # If still in command mode (e.g. error)
                 self.command_buffer = "" # Clear buffer for next command
            # If _execute_command returned False (quit), this won't matter
        elif 32 <= key <= 126: # Printable characters
            self.command_buffer += chr(key)
        return True # Default continue running, _execute_command handles quit

    def _execute_command(self):
        parts = self.command_buffer.strip().split()
        command = parts[0] if parts else ""
        args = parts[1:]

        self.status_message = "" # Clear previous command status

        if command == 'q':
            if self.dirty:
                self.status_message = "E37: No write since last change (add ! to override)"
                self.mode = NORMAL_MODE # Stay in editor, but clear command buffer
                self.command_buffer = ""
                return True
            return False # Signal to quit main loop
        elif command == 'q!':
            return False # Signal to quit main loop
        elif command == 'w':
            filename_to_save = args[0] if args else None
            if self._save_file(filename_to_save):
                 pass # Success message is set by _save_file
            else:
                 pass # Error message set by _save_file
            self.mode = NORMAL_MODE
        elif command == 'wq':
            # Try to save, if successful then quit
            filename_to_save = args[0] if args else None # wq can also take a filename
            if self._save_file(filename_to_save):
                return False # Quit after successful save
            else:
                # Save failed, stay in command mode or normal mode? Vim goes to normal.
                self.status_message += " (wq failed)"
                self.mode = NORMAL_MODE
        elif command == 'e': # Basic :e <filename>
            if len(args) > 0:
                if self.dirty:
                    self.status_message = "E37: No write since last change. Use :e! to discard."
                    self.mode = NORMAL_MODE
                    return True
                self.filename = args[0]
                self._load_file(self.filename)
                self.cursor_x = 0
                self.cursor_y = 0
                self.top_row = 0
                self.dirty = False
                self.status_message = f"Loaded \"{self.filename}\""
            else:
                self.status_message = "E499: Missing filename for :e"
            self.mode = NORMAL_MODE
        elif command == 'e!':
             if len(args) > 0:
                self.filename = args[0]
                self._load_file(self.filename)
                self.cursor_x = 0
                self.cursor_y = 0
                self.top_row = 0
                self.dirty = False
                self.status_message = f"Loaded \"{self.filename}\" (changes discarded)"
             else:
                self.status_message = "E499: Missing filename for :e!"
             self.mode = NORMAL_MODE

        else:
            self.status_message = f"Not an editor command: {self.command_buffer}"
            self.mode = NORMAL_MODE # Go back to normal mode on unknown command
        
        self.command_buffer = "" # Clear command buffer after execution attempt
        return True


    def run(self):
        # Initial setup for curses
        curses.noecho()  # Don't echo pressed keys to screen
        curses.cbreak()   # React to keys instantly, without Enter
        self.stdscr.keypad(True) # Enable special keys (arrows, F-keys, etc.)

        # Basic color support (optional, but good for line numbers/tilde)
        try:
            curses.start_color()
            curses.use_default_colors() # Allow use of terminal's default background
            # Pair 1: Foreground for line numbers, tildes (e.g., blue on default bg)
            curses.init_pair(1, curses.COLOR_BLUE, -1) # -1 for default background
            curses.init_pair(2, curses.COLOR_GREEN, -1) # For status messages
        except curses.error:
            # Color setup failed (e.g., terminal doesn't support colors)
            pass


        while True:
            self._ensure_cursor_bounds()
            self._scroll_if_needed() # Adjust viewport before rendering
            self._render()

            current_mode_cursor_pos_y = self.cursor_y - self.top_row
            current_mode_cursor_pos_x = (self.cursor_x - self.left_col) + 4 # For line numbers

            if self.mode == COMMAND_MODE:
                # For command mode, cursor is on the command line
                self.stdscr.move(self.height -1, len(self.command_buffer) + 1) # +1 for ':'
            elif 0 <= current_mode_cursor_pos_y < (self.height - 2) and \
                 4 <= current_mode_cursor_pos_x < self.width :
                try: # Defensive move
                    self.stdscr.move(current_mode_cursor_pos_y, current_mode_cursor_pos_x)
                except:
                    pass # Ignore if move fails for some reason
            
            key = self.stdscr.getch() # Get user input (blocking)

            continue_running = True
            if self.mode == NORMAL_MODE:
                continue_running = self._handle_normal_mode_input(key)
            elif self.mode == INSERT_MODE:
                continue_running = self._handle_insert_mode_input(key)
            elif self.mode == COMMAND_MODE:
                continue_running = self._handle_command_mode_input(key)
            
            if not continue_running:
                break # Exit main loop

def main(stdscr_arg):
    # Check for filename argument
    filename_arg = None
    if len(sys.argv) > 1:
        filename_arg = sys.argv[1]

    editor = PyVim(stdscr_arg, filename_arg)
    editor.run()

if __name__ == '__main__':
    # curses.wrapper handles terminal setup and teardown (restore on exit/error)
    curses.wrapper(main)