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
    color: str = 'green'

    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 add_tree_if_possible(self, min_distance_2, x, y, **kparams):
        _nearest_tree, distance_2 = self.kd_tree.search_nn((x, y))
        if distance_2 > min_distance_2:
            self.kd_tree.add((x, y))
            self.trees.append(Tree(x, y, **kparams))

    @property
    def xs_ys(self):
        xs, ys = [], []
        for tree in self.trees:
            xs.append(tree.x)
            ys.append(tree.y)
        return xs, ys

    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])