tree.py

from dataclasses import dataclass
from collections import UserList
import kdtree


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

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

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

    def __str__(self):
        return f"{self.type} ({self.description}), {self.radius or '?'} m (X={self.x:.1f}, Y={self.y:.1f})"


class Forest(UserList):
    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.data = 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.append(Tree(x, y, **kparams))

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

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

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