Quick page table emulation in python

vm.py

class VM:
    def __init__(self, page_size):
        self.page_size = page_size - 1
        self.page_size_bits = self.page_size.bit_length()
        self.page_table = {}
        self.tlb = {}
    
    def translate(self, address) -> int:
        try:
            return self.tlb[address]
        except KeyError:
            pass
        page_offset = address & self.page_size
        page_location = address >> self.page_size_bits
        try:
            phys_page = self.page_table[page_location]
            physical_location = phys_page + page_offset
            self.tlb[address] = physical_location
            return physical_location
        except KeyError:
            print(f'Address {hex(address)} not in VA Table. (page: {hex(page_location)})')
            for i in self.page_table:
                print(hex(i))
            raise AssertionError
        
    def map_pages(self, physical_addr, virtual_addr, size):
        for i in range(size // self.page_size):
            self.page_table[virtual_addr + (i * (self.page_size + 1)) >> self.page_size_bits] = physical_addr + (i * (self.page_size + 1))
            
   
vm = VM(page_size = 0x4000)
vm.map_pages(0x0, 0x100000000, 0x4000)
vm.map_pages(0x8000, 0x100004000, 0x8000)
vm.map_pages(0x20000, 0x10000c000, 0x28000)

r_base = 0x10000c000
import random
lst = [r_base+random.randrange(0x0, 0x10000) for i in range(1000000)]
for i in range(1000000):
    vm.translate(lst[i])