highlight_problem.py

"""
filename: highlight_problem.py
author: Tristan Hilbert
desc: This was a problem I brought up to a research fellow.
  I was trying to show that storing highlighted text that can
  have multiple layers is difficult. Below I have created a python
  function and a series of tests. Try to get the tests to pass by
  modifying the function below. Feel free to add functions but do not modify the tests!
"""

from dataclasses import dataclass
from enum import Enum
import time
from typing import Optional

class Input:
    pass # see below
class Command:
    pass # see below
class Label:
    pass # see below


"""
---------------------------------------------------------------------------------------
START HERE
---------------------------------------------------------------------------------------
"""

class Highlights:
    pass # or whatever you would like to call it.

def highlight_datastructure(input: Input, structure: Optional[Highlights]) -> Highlights | list[Label]:
    pass


"""
---------------------------------------------------------------------------------------
PLEASE DO NOT MODIFY BELOW
---------------------------------------------------------------------------------------
"""

# can handle a simple addition without problem
class Command(Enum):
    CREATE = 0 # Create an empty data structure reference for the caller
    ADD = 1 # Provides a start and end and a new label, add input to datastructure
    CHECK = 2 # Asks for a check of the values highlighting start to end
    REMOVE = 3 # Remove the label(s) attached to the characters from start to end

class Label(Enum):
    A = 1
    B = 2
    C = 3
    D = 4
    E = 5

@dataclass
class Input:
    corpus: str # the overall string being highlighted, feel free to cache this value
    command: Command # see above
    start: int # inclusive range, all characters including starting index should be highlighted
    end: int # inclusive range, all characters including ending index should be highlighted
    label: Optional[Label] # Label that the datastructure should keep track of

# Make sure it can handle additions
def test_0():
    corpus = "I am Lorem Ipsum"
    length = len(corpus)
    structure = highlight_datastructure(Input(corpus, Command.CREATE, 0, 0, None), None)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, 4, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, 5, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, 7, Label.A), structure)

# Single Highlight works well
def test_1():
    corpus = "I am Lorem Ipsum I am Lorem Ipsum"
    length = len(corpus)
    structure = highlight_datastructure(Input(corpus, Command.CREATE, 0, 0, None), None)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, length - 1, Label.A), structure)
    assert highlight_datastructure(Input(corpus, Command.CHECK, 0, length - 1)) == [Label.A]
    assert highlight_datastructure(Input(corpus, Command.CHECK, 0, 1)) == [Label.A]
    assert highlight_datastructure(Input(corpus, Command.CHECK, 1, 2)) == [Label.A]

# Distinct Highlights work
def test_2():
    corpus = "I am Lorem Ipsum I am Lorem Ipsum"
    length = len(corpus)
    structure = highlight_datastructure(Input(corpus, Command.CREATE, 0, 0, None), None)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, 2, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 3, 5, Label.B), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 8, 10, Label.C), structure)
    assert highlight_datastructure(Input(corpus, Command.CHECK, 0, 0)) == [Label.A]
    assert highlight_datastructure(Input(corpus, Command.CHECK, 4, 4)) == [Label.B]
    assert highlight_datastructure(Input(corpus, Command.CHECK, 9, 9)) == [Label.C]

# Checks on Highlights work
def test_3():
    corpus = "I am Lorem Ipsum I am Lorem Ipsum"
    length = len(corpus)
    structure = highlight_datastructure(Input(corpus, Command.CREATE, 0, 0, None), None)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, 2, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 3, 5, Label.B), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 8, 10, Label.C), structure)
    assert sorted(highlight_datastructure(Input(corpus, Command.CHECK, 0, length - 1))) == [Label.A, Label.B, Label.C]

# Mid Removal Works
def test_4():
    corpus = "I am Lorem Ipsum I am Lorem Ipsum"
    length = len(corpus)
    structure = highlight_datastructure(Input(corpus, Command.CREATE, 0, 0, None), None)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, length - 1, Label.A), structure)
    assert highlight_datastructure(Input(corpus, Command.CHECK, 0, 0)) == [Label.A]
    assert highlight_datastructure(Input(corpus, Command.CHECK, 8, 8)) == [Label.A]
    structure = highlight_datastructure(Input(corpus, Command.REMOVE, 6, 10, Label.A), structure)
    assert highlight_datastructure(Input(corpus, Command.CHECK, 8, 8)) == []
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, length - 1, Label.A), structure)
    assert highlight_datastructure(Input(corpus, Command.CHECK, 8, 8)) == [Label.A]
    structure = highlight_datastructure(Input(corpus, Command.REMOVE, 10, length - 1, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.REMOVE, 6, length - 1, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 0, length - 1, Label.B), structure)
    assert highlight_datastructure(Input(corpus, Command.CHECK, 8, 8)) == [Label.B]

# Board Example
def test_5():
    corpus = "I am Lorem Ipsum I am Lorem Ipsum"
    length = len(corpus)
    structure = highlight_datastructure(Input(corpus, Command.CREATE, 0, 0, None), None)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 3, 5, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 8, length - 1, Label.B), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 11, 20, Label.C), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 12, 14, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.ADD, 15, 25, Label.C), structure)
    assert sorted(highlight_datastructure(Input(corpus, Command.CHECK, 0, length - 1))) == [Label.A, Label.B, Label.C]
    assert highlight_datastructure(Input(corpus, Command.CHECK, 4, 4)) == [Label.A]
    assert sorted(highlight_datastructure(Input(corpus, Command.CHECK, 9, 20))) == [Label.A, Label.B, Label.C]
    assert sorted(highlight_datastructure(Input(corpus, Command.CHECK, 21, 24))) == [Label.A, Label.B]

    structure = highlight_datastructure(Input(corpus, Command.REMOVE, 0, length - 1, Label.A), structure)
    structure = highlight_datastructure(Input(corpus, Command.REMOVE, 0, length - 1, Label.B), structure)
    structure = highlight_datastructure(Input(corpus, Command.REMOVE, 0, length - 1, Label.C), structure)
    assert highlight_datastructure(Input(corpus, Command.CHECK, 0, length - 1)) == []


# Don't be slow
def test_6():
    corpus = "I am Lorem Ipsum I am Lorem Ipsum"
    length = len(corpus)
    structure = highlight_datastructure(Input(corpus, Command.CREATE, 0, 0, None), None)
    start = time.time()
    for i in range(10000):    
        structure = highlight_datastructure(Input(corpus, Command.ADD, 3, 5, Label.A), structure)
        structure = highlight_datastructure(Input(corpus, Command.ADD, 8, length - 1, Label.B), structure)
        structure = highlight_datastructure(Input(corpus, Command.ADD, 11, 20, Label.C), structure)
        structure = highlight_datastructure(Input(corpus, Command.ADD, 12, 14, Label.A), structure)
        structure = highlight_datastructure(Input(corpus, Command.ADD, 15, 25, Label.C), structure)
        structure = highlight_datastructure(Input(corpus, Command.REMOVE, 0, length - 1, Label.B), structure)

    assert sorted(highlight_datastructure(Input(corpus, Command.CHECK, 21, 24))) == [Label.A]
    assert sorted(highlight_datastructure(Input(corpus, Command.CHECK, 11, 11))) == [Label.C]
    structure = highlight_datastructure(Input(corpus, Command.REMOVE, 0, length - 1, Label.C), structure)
    assert highlight_datastructure(Input(corpus, Command.CHECK, 0, length - 1)) == [Label.A]

    assert time.time() - start < 300 # if ur taking more than 5 minutes ur doing something wrong

def main():
    tests = [
        test_0,
        test_1,
        test_2,
        test_3,
        test_4,
        test_5,
        test_6
    ]

    for test in tests:
        test()
        print(f"Test {test.__name__} passed!")

if __name__ == "__main__":
    main()