tree.py

from dataclasses import dataclass
import kdtree

@dataclass
class Tree:
    x: float
    y: float
    z: float = 0
    height: float = None
    radius: float = None
    description: str = None
    type: str = None

    def __len__(self):
        return 2 # x & y

    def __getitem__(self, i):
        return [self.x, self.y][i]


class Forest:
    def __init__(self, existing_trees=[]):
        if existing_trees:
            self.kd_tree = kdtree.create(existing_trees, dimensions=2)
        else:
            self.kd_tree = kdtree.create([(0, 0)], dimensions=2)

        self.trees = existing_trees

    def __str__(self):
        return f"Forest with {len(self.trees)} trees."

    def __repr__(self):
        return "\n".join([str(self)] + [str(tree) for tree in self.trees])