from selenium.webdriver.remote.webdriver import WebDriver

class LocalStorage():
	# from https://stackoverflow.com/a/46361900
	def __init__(self, driver: WebDriver):
		self.driver = driver

	def __len__(self):
		return self.driver.execute_script("return window.localStorage.length;")

	def items(self):
		return self.driver.execute_script( \
			"var ls = window.localStorage, items = {}; " \
			"for (var i = 0, k; i < ls.length; ++i) " \
			"	items[k = ls.key(i)] = ls.getItem(k); " \
			"return items;")

	def keys(self):
		return self.driver.execute_script("return Object.keys(window.localStorage);")

	def get(self, key):
		return self.driver.execute_script("return window.localStorage.getItem(arguments[0]);", key)

	def set(self, key, value):
		self.driver.execute_script("window.localStorage.setItem(arguments[0], arguments[1]);", key, value)

	def has(self, key):
		return key in self.keys()

	def remove(self, key):
		self.driver.execute_script("window.localStorage.removeItem(arguments[0]);", key)

	def clear(self):
		self.driver.execute_script("window.localStorage.clear();")

	def __getitem__(self, key) :
		value = self.get(key)
		if value is None :
			raise KeyError(key)
		return value

	def __setitem__(self, key, value):
		self.set(key, value)

	def __contains__(self, key):
		return key in self.keys()

	def __iter__(self):
		return self.items().__iter__()

	def __repr__(self):
		return self.items().__str__()
