from dataclasses import dataclass, field
from typing import TYPE_CHECKING
import numpy as np
import pymunk
import pymunk.pygame_util
from symaware.base.models import Entity as BaseEntity
from symaware.base.models import NullDynamicalModel
if TYPE_CHECKING:
# String type hinting to support python 3.9
import pygame
# Forwards reference
from .dynamical_model import DynamicalModel
[docs]
@dataclass(frozen=True)
class Entity(BaseEntity):
"""
Abstract class for the entities using the Pymunk physics engine.
Args
----
model:
Dynamical model of the entity. Must be a subclass of :class:`.PymunkDynamicalModel`
position:
Initial position of the entity
angle:
Initial angle of the entity
body:
Pymunk body of the entity. It is automatically created
shape:
Pymunk shape of the entity. It is automatically created
color:
Color of the entity used for visualisation
mass:
Mass of the entity
friction:
Friction of the entity
elasticity:
Elasticity of the entity
"""
model: "DynamicalModel" = field(default_factory=NullDynamicalModel)
position: np.ndarray = field(default_factory=lambda: np.zeros(2))
angle: float = field(default=0.0)
body: "pymunk.Body" = field(default_factory=pymunk.Body)
color: "pygame.Color | None" = field(default=None)
mass: float = field(default=1)
friction: float = field(default=0.7)
elasticity: float = field(default=0)
_shape: "pymunk.Shape" = field(default=None)
def __post_init__(self):
object.__setattr__(
self, "position", np.pad(self.position[:2], (0, max(0, 2 - len(self.position))), mode="constant")
)
self.body.position = pymunk.Vec2d(self.position[0], self.position[1])
self.body.angle = self.angle
self.body.velocity = pymunk.Vec2d(0, 0)
self._shape.mass = self.mass
self._shape.friction = self.friction
self._shape.elasticity = self.elasticity
if self.color is not None:
self._shape.color = self.color
[docs]
def initialise(self, space: pymunk.Space):
space.add(self.body, self._shape)
if not isinstance(self.model, NullDynamicalModel):
self.model.initialise(self)
@property
def shape(self) -> "pymunk.Shape":
return self._shape
def __hash__(self) -> int:
return hash((super().__hash__(), id(self)))
[docs]
@dataclass(frozen=True)
class SphereEntity(Entity):
radius: float = field(default=1.0)
def __post_init__(self):
object.__setattr__(self, "_shape", pymunk.Circle(self.body, radius=self.radius))
super().__post_init__()
def __hash__(self) -> int:
return hash((super().__hash__(), id(self)))
[docs]
@dataclass(frozen=True)
class BoxEntity(Entity):
sizes: np.ndarray = field(default_factory=lambda: np.ones(2))
def __post_init__(self):
object.__setattr__(self, "_shape", pymunk.Poly.create_box(self.body, size=(self.sizes[0], self.sizes[1])))
object.__setattr__(self, "sizes", np.pad(self.sizes[:2], (0, max(0, 2 - len(self.sizes))), mode="constant"))
super().__post_init__()
def __hash__(self) -> int:
return hash((super().__hash__(), id(self)))