Python - Data Structures

linked_list.py

import unittest

class Node():
	def __init__(self, data):
		self.data = data
		self.next = None

	def __str__(self):
		return f"Node(data={self.data}, next={str(self.next)})"

class LinkedList():
	def __init__(self, head=None):
		self.head = head

	def append(self, data):
		node = Node(data=data)
		if self.head == None:
			self.head = node

		temp = self.head
		while temp.next != None:
			temp = temp.next

		temp.next = node

	# insert first
	def prepend(self, data):
		node = Node(data=data)
		if self.head == None:
			self.head = node

		node.next = self.head
		self.head = node

	# insert after specific data
	def insert_at(self, data, prev_data):
		node = Node(data=data)

		if self.head == None:
			self.head = node
		
		temp = self.head
		while temp != None:
			if temp.data == prev_data:
				next_node = temp.next
				node.next = next_node
				temp.next = node
				break

			temp = temp.next

	# remove first
	def pop(self):
		pass

	# remove last
	def remove(self):
		pass

	# remove at specific data
	def remove_at(self, data):
		pass

	# find value
	def search(self, data):
		pass

	def __str__(self):
		# print the linked list into following format
		# 1 -> 5 -> 3 -> None
		sequence = ""
		temp = self.head
		while temp != None:
			sequence += f"{temp.data} -> "
			if temp.next == None:
				sequence += f"None"

			temp = temp.next

		return sequence

# Unit Test
class TestLinkedList(unittest.TestCase):
	def setUp(self):
		self.linked_list = LinkedList()
		head = Node(data=1)
		self.linked_list.head = head

	def test_append(self):
		self.linked_list.append(20)

		self.assertEqual(str(self.linked_list), "1 -> 20 -> None")

	def test_prepend(self):
		self.linked_list.prepend(3)

		self.assertEqual(str(self.linked_list), "3 -> 1 -> None")

	def test_insert_at(self):
		self.linked_list.insert_at(4, 1)
		self.assertEqual(str(self.linked_list), "1 -> 4 -> None")
		
		self.linked_list.insert_at(2, 1)
		self.assertEqual(str(self.linked_list), "1 -> 2 -> 4 -> None")

if __name__ == '__main__':
	unittest.main()