pygame smoke effect

smoke.py

# ispired by: https://www.youtube.com/watch?v=n-tLlBNgS3g
import random
import sys
from contextlib import contextmanager

import pygame

Vec2 = pygame.math.Vector2

FPS = 60
GREY = (128, 128, 128)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)


def smooth_stop(t, power):
    def flip(x):
        return 1 - x

    def dumb_pow(x, p):
        res = 1.0
        while p > 0:
            res *= x
            p -= 1
        return res

    return flip(dumb_pow(flip(t), power))


@contextmanager
def managed_pygame(*args, **kwargs):
    width = kwargs.get("width", 500)
    height = kwargs.get("height", 500)
    caption = kwargs.get("caption", "pygame window")

    pygame.init()
    screen = pygame.display.set_mode((width, height))
    pygame.display.set_caption(caption)

    clock = pygame.time.Clock()

    try:
        yield screen, clock
    finally:
        pygame.quit()
        sys.exit()


class Mover:
    def __init__(self) -> None:
        self.position = Vec2(0, 0)
        self.velocity = Vec2(0, 0)
        self.acceleration = Vec2(0, 0)

    def apply_force(self, force: Vec2) -> None:
        self.acceleration += force

    def update(self) -> None:
        self.velocity += self.acceleration
        self.position += self.velocity
        self.acceleration = Vec2(0, 0)


SMOKE = pygame.image.load("/Users/tomowarkar/Downloads/smoke.png")


class Bullet(pygame.sprite.Sprite, Mover):
    def __init__(self):
        super().__init__()

        self.surf = SMOKE
        self.surf = pygame.transform.scale(self.surf, (30, 30))

        self.rect = self.surf.get_rect()

        Mover.__init__(self)

        self.lifetime = 300
        self.angle = 0
        self.scale = 1

    def update(self):
        self.apply_force(self.velocity * -0.02)
        Mover.update(self)

        self.rect.center = self.position
        self.angle += self.velocity.magnitude()

        self.scale += 10 / 300

        if self.lifetime < 0:
            self.kill()
        else:
            self.lifetime -= 1

    def draw(self, surface: pygame.Surface):
        rotated = pygame.transform.rotozoom(self.surf, self.angle, self.scale)
        rect = rotated.get_rect(center=self.rect.center)
        surface.blit(rotated, rect)


class Cannon:
    def __init__(self) -> None:
        self.position = Vec2(100, 250)
        self.direction = Vec2(1, 0)
        self.length = 20
        self.color = BLACK
        self.width = 1

    @property
    def dest(self):
        return self.direction * self.length + self.position

    def update(self):
        mouse = pygame.mouse.get_pos()
        dir = Vec2(mouse) - self.position
        self.direction = dir.normalize()

    def draw(self, surface: pygame.Surface):
        pygame.draw.line(surface, self.color, self.position, self.dest, self.width)


with managed_pygame() as (screen, clock):
    pygame.key.set_repeat(100)

    c = Cannon()
    bullets = pygame.sprite.Group()

    font = pygame.font.SysFont(None, 20)

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    b = Bullet()
                    b.position = c.dest
                    b.apply_force(c.direction * (4 + random.random()))
                    bullets.add(b)

        screen.fill(WHITE)

        c.update()
        c.draw(screen)

        for bullet in bullets:
            bullet.update()
            bullet.draw(screen)

        clock.tick(FPS)
        screen.blit(
            font.render(
                "FPS: %.2f, bullets: %d" % (clock.get_fps(), len(bullets)),
                True,
                BLACK,
            ),
            (0, 0),
        )
        pygame.display.update()