jsonrep.mojo

from collections import List, Dict
from utils.variant import Variant

@value
struct JsonNode(CollectionElement, Stringable):
    var value: Pointer[Variant[NoneType, Bool, Float64, String, List[Self], Dict[String, Self]]]

    fn __init__(inout self: Self, value: Variant[NoneType, Bool, Float64, String, List[Self], Dict[String, Self]]):
        self.value = Pointer[Variant[NoneType, Bool, Float64, String, List[Self], Dict[String, Self]]].alloc(1)
        self.value[0] = value

    # Silly specializations to handle the different types
    fn __init__(inout self: Self, value: Dict[String, Bool]) raises:
        self.value = Pointer[Variant[NoneType, Bool, Float64, String, List[Self], Dict[String, Self]]].alloc(1)
        var dict = Dict[String, Self]()
        for key in value.keys():
            dict[key[]] = JsonNode(value=value[key[]])
        self.value[0] = dict

    fn __init__(inout self: Self, value: Dict[String, Float64]) raises:
        self.value = Pointer[Variant[NoneType, Bool, Float64, String, List[Self], Dict[String, Self]]].alloc(1)
        var dict = Dict[String, Self]()
        for key in value.keys():
            dict[key[]] = JsonNode(value=value[key[]])
        self.value[0] = dict

    fn __init__(inout self: Self, value: Dict[String, String]) raises:
        self.value = Pointer[Variant[NoneType, Bool, Float64, String, List[Self], Dict[String, Self]]].alloc(1)
        var dict = Dict[String, Self]()
        for key in value.keys():
            dict[key[]] = JsonNode(value=String(value[key[]]))
        self.value[0] = dict

    fn __copyinit__(inout self: Self, existing: Self):
        self.value = Pointer[Variant[NoneType, Bool, Float64, String, List[Self], Dict[String, Self]]].alloc(1)
        var other_val = existing.value[0]
        self.value[0] = other_val

    fn __moveinit__(inout self, owned existing: Self):
        self.value = existing.value
        existing.value = Pointer[Variant[NoneType, Bool, Float64, String, List[Self], Dict[String, Self]]].alloc(1)

    fn __del__(owned self):
        return self.value.free()

    fn __str__(self) -> String:
        try:
            if self.value[0].isa[NoneType]():
                return "null"
            elif self.value[0].isa[Bool]():
                var bool_val = self.value[0].get[Bool]()[]
                if bool_val:
                    return "true"
                else:
                    return "false"
            elif self.value[0].isa[Float64]():
                return str(self.value[0].get[Float64]()[])
            elif self.value[0].isa[String]():
                return "\"" + self.value[0].get[String]()[] + "\""
            elif self.value[0].isa[List[Self]]():
                var list = self.value[0].get[List[Self]]()[]
                if len(list) == 0:
                    return "[]"
                var stringBuf: String = "[\n"
                var first = True
                for i in range(len(list)):
                    if not first:
                        stringBuf += ",\n"
                    first = False
                    var string_val = String(list[i])
                    var lines = string_val.split("\n")
                    for x in range(len(lines)):
                        stringBuf += "  " + lines[x]
                        if x < len(lines) - 1:
                            stringBuf += "\n"
                return stringBuf + "\n]"
            elif self.value[0].isa[Dict[String, Self]]():
                var dict = self.value[0].get[Dict[String, Self]]()[]
                var str_list = List[String]()
                if len(dict) == 0:
                    return "{}"
                var stringBuf: String = "{\n"
                var first = True
                for key in dict.keys():
                    if not first:
                        stringBuf += ",\n"
                    first = False
                    stringBuf += "  \"" + key[] + "\": " 
                    var string_val = String(dict[key[]])
                    var lines = string_val.split("\n")
                    for x in range(len(lines)):
                        if x > 0:
                            stringBuf += "  "
                        stringBuf += lines[x]
                        if x < len(lines) - 1:
                            stringBuf += "\n"
                return stringBuf + "\n}"
            else:
                return "unknown"
        except:
            return "error"
    
    fn __repr__(self) -> String:
        return self.__str__()

    fn get[T:CollectionElement](self) -> T:
        return self.value[0].get[T]()[]
    
    fn __getitem__(self, key: String) raises -> JsonNode:
        if self.value[0].isa[Dict[String, Self]]():
            return self.value[0].get[Dict[String, Self]]()[][key]
        else:
            raise "Not a dictionary"

    fn __getitem__(self, index: Int) raises -> JsonNode:
        if self.value[0].isa[List[Self]]():
            return self.value[0].get[List[Self]]()[][index]
        else:
            raise "Not a list"

fn load() raises -> JsonNode:
    var data = Dict[String, JsonNode]()
    data["name"] = JsonNode(value=String("John"))
    data["isStudent"] = JsonNode(value=True)
    data["age"] = JsonNode(value=Float64(25))
    var friends = List[JsonNode]()
    friends.append(JsonNode(value=String("Alice")))
    friends.append(JsonNode(value=String("Bob")))
    data["friends"] = JsonNode(value=friends)
    return JsonNode(value=data)


fn load_dict() raises -> JsonNode:
    var data = Dict[String, String]()
    data["name"] = "John"
    data["email"] = "fake@gmail.com"
    data["age"] = "25"

    return JsonNode(value=data)

fn main() raises:
    var data = load()
    print(data)
    var name = data.get[Dict[String, JsonNode]]()["name"].get[String]()
    print(name)
    var friend1 = data.get[Dict[String, JsonNode]]()["friends"].get[List[JsonNode]]()[0].get[String]()
    print(friend1)
    var friend2 = data["friends"][1].get[String]()
    print(friend2)

    var data2 = load_dict()
    print(data2["name"].get[String]())
    print(data2)