Reading from stdin using mojo and recreating the Python input function

stdin.mojo

from builtin.io import _dup
from memory import UnsafePointer, memcpy
from sys import os_is_windows, external_call


@value
struct stdin:
    """A read only file handle to the stdin stream."""

    alias file_descriptor = 0
    alias mode = "r"
    var handle: UnsafePointer[NoneType]
    """The file handle to the stdin stream."""

    fn __init__(inout self):
        """Creates a file handle to the stdin stream."""
        var handle: UnsafePointer[NoneType]

        @parameter
        if os_is_windows():
            handle = external_call["_fdopen", UnsafePointer[NoneType]](
                _dup(Self.file_descriptor), Self.mode.unsafe_ptr()
            )
        else:
            handle = external_call["fdopen", UnsafePointer[NoneType]](
                _dup(Self.file_descriptor), Self.mode.unsafe_ptr()
            )
        self.handle = handle

    fn readline(self) -> String:
        """Reads an entire line from stdin or until EOF. Lines are delimited by a newline character.
        
        Returns:
            The line read from the stdin.
        
        Examples:

        ```mojo
        from sys.std import stdin
        var line = stdin().readline()
        print(line)
        ```
        Assuming the above program is named `my_program.mojo`, feeding it `Hello, World` via stdin would output:

        ```bash
        echo "Hello, World" | mojo run my_program.mojo
        Hello, World # Output from print
        ```
        
        The program can also be run interactively by typing `Hello, World` and then `Enter` to input the terminating newline.
        ```bash
        mojo run my_program.mojo
        Hello, World # User input via the terminal
        Hello, World # Output from print
        ```.
        """
        return self.read_until_delimiter("\n")

    fn read_until_delimiter(self, delimiter: String) -> String:
        """Reads an entire line from stdin, which is delimited by `delimiter`. 
        Does not include the delimiter in the result.

        Args:
            delimiter: The delimiter to read until.
        
        Returns:
            The line read from the stdin.
        
        Examples:

        ```mojo
        from sys.std import stdin

        fn main():
            var line = stdin().read_until_delimiter(",")
            print(line)
        ```

        Assuming the above program is named `my_program.mojo`, feeding it `Hello, World` via stdin would output:

        ```bash
        echo "Hello, World" | mojo run my_program.mojo
        Hello # Output from print
        ```
        
        The program can also be run interactively by typing `Hello, World` and then `Enter` to input the terminating newline.
        ```bash
        mojo run my_program.mojo
        Hello, World # User input via the terminal
        Hello # Output from print
        ```.
        """
        var buffer = UnsafePointer[UInt8].alloc(1)
        var bytes_read = external_call[
            "getdelim",
            Int,
            UnsafePointer[UnsafePointer[UInt8]],
            UnsafePointer[UInt32],
            Int,
            UnsafePointer[NoneType],
        ](
            UnsafePointer[UnsafePointer[UInt8]].address_of(buffer),
            UnsafePointer[UInt32].address_of(UInt32(1)),
            ord(delimiter),
            self.handle,
        )
        buffer[bytes_read - 1] = 0
        return String(buffer, bytes_read)

    fn read(self, size: Int = -1) -> String:
        """Reads up to `size` characters from stdin or until EOF.
        Only supports ASCII characters that can be represented as a single byte.

        Args:
            size: The number of characters to read from the file handle. If not provided, reads until EOF.
        
        Returns:
            The characters read from the stdin.
        
        Examples:

        ```mojo
        from sys.std import stdin

        fn main():
            var line = stdin().read(5)
            print(line)
        ```

        Assuming the above program is named `my_program.mojo`, feeding it `Hello, World` via stdin would output:

        ```bash
        echo "Hello, World" | mojo run my_program.mojo
        Hello # Output from print
        ```
        
        The program can also be run interactively by typing `Hello, World` and then `Enter` to input the terminating newline.
        ```bash
        mojo run my_program.mojo
        Hello, World # User input via the terminal
        Hello # Output from print
        ```.
        """
        if size == -1:
            return self.read_until_eof()

        var buffer = UnsafePointer[UInt8].alloc(size + 1)
        var bytes_read = self.read(buffer, size)
        buffer[bytes_read] = 0
        return String(buffer, bytes_read + 1)

    fn read(self, buffer: UnsafePointer[UInt8], size: Int) -> Int:
        """Reads up to `size` bytes from the file handle into the buffer or until EOF.
        Only supports ASCII characters that can be represented as a single byte.

        Args:
            buffer: The buffer to read the data into.
            size: The number of bytes to read from the file handle.

        Returns:
            The number of bytes read from stdin.
        """
        var total_bytes_read = 0
        while (
            external_call["feof", Int, UnsafePointer[NoneType]](self.handle)
            == 0
            and total_bytes_read < size
        ):
            var bytes_read = external_call[
                "fread",
                Int,
                UnsafePointer[UInt8],
                Int,
                Int,
                UnsafePointer[NoneType],
            ](buffer + total_bytes_read, sizeof[UInt8](), size, self.handle)

            # TODO: Remove EOF character from bytes read to remove the trailing newline that happens otherwise.
            # Should the newline that submits the input to stdin be considered part of the input? 
            if (
                external_call["feof", Int, UnsafePointer[NoneType]](self.handle)
                != 0
            ):
                bytes_read -= 1

                # TODO: Capture error and raise with ferror
            total_bytes_read += bytes_read

        return total_bytes_read

    fn read_until_eof(self) -> String:
        var buffer_size = 8192
        var buffer = UnsafePointer[UInt8].alloc(buffer_size)
        var total_bytes_read = 0
        while (
            external_call["feof", Int, UnsafePointer[NoneType]](self.handle)
            == 0
        ):
            var bytes_read = external_call[
                "fread",
                Int,
                UnsafePointer[UInt8],
                Int,
                Int,
                UnsafePointer[NoneType],
            ](
                buffer + total_bytes_read,
                sizeof[UInt8](),
                buffer_size,
                self.handle,
            )
            if (
                external_call["feof", Int, UnsafePointer[NoneType]](self.handle)
                != 0
            ):
                bytes_read -= 1
            total_bytes_read += bytes_read

            if bytes_read == buffer_size:
                var new = UnsafePointer[UInt8].alloc(bytes_read * 2)
                buffer_size = bytes_read * 2
                memcpy(new, buffer, total_bytes_read)
                buffer.free()
                buffer = new

        buffer[total_bytes_read] = 0
        return String(buffer, total_bytes_read + 1)

    fn close(self):
        _ = external_call["fclose", Int32](self.handle)

    fn __del__(owned self):
        self.close()

    fn __enter__(self) -> Self:
        return self

    fn __exit__(self):
        """Closes the file handle."""
        self.close()



fn input(prompt: String = "") -> String:
    if prompt != "":
        print(prompt, end="")
    return stdin().readline()


fn main():
    var user_input = input("What's your name?")
    print(user_input)
    print("Reading a line from stdin")
    print(stdin().readline())

    print("Reading until ','")
    print(stdin().read_until_delimiter(","))

    print("Reading 8 bytes into a pointer from stdin")
    var buf = UnsafePointer[UInt8].alloc(16)
    var bytes_read = stdin().read(buf, 8)
    buf[bytes_read] = 0
    print(String(buf, bytes_read + 1))

    print("Reading 5 bytes from stdin")
    print(stdin().read(5))

    print("Reading until EOF from stdin, ctrl+D to send EOF")
    print(stdin().read_until_eof())

    print("Reading until EOF from stdin, ctrl+D to send EOF")
    print(stdin().read())