どう書く？ ピラミッドを作る

ピラミッドを作る どう書く？org

帰りの電車で見つけて楽しそうだったので書いてみた。

ピラミッドなので、一番下の階（？）からコツコツ建てていく感じです。

# pyramid.py
# coding: utf-8

"""num階建てのピラミッドを返す"""
def pyramid(num):
  return build_pyramid(ground(num))

"""num階建てピラミッドの一階を返す"""
def ground(num):
  return "*" * (num * 2 - 1)

"""下の階を受け取って一段上の階を返す"""
def next_floor(floor):
  return floor.replace("*", " ", 1)[:-1]

"""下の階を受け取って、その上に建つピラミッドを返す"""
def build_pyramid(under_floor):
  current_floor = next_floor(under_floor)
  if current_floor.strip() == "":
    return under_floor
  return build_pyramid(current_floor) + "\n" + under_floor

if __name__ == '__main__':
  import sys
  print pyramid(int(sys.argv[1]))

テスト

# test_pyramid.py
# coding: utf-8

from pyramid import *
import unittest

class TestPyramid(unittest.TestCase):
  def test_4_floor(self):
    self.assertEquals(_4_floor_pyramid(), pyramid(4))

  def test_3_floor(self):
    self.assertEquals(_3_floor_pyramid(), pyramid(3))

  def test_ground_of_4_floor(self):
    self.assertEquals("*******", ground(4))

  def test_ground_of_3_floor(self):
    self.assertEquals("*****", ground(3))

  def test_ground_of_0_floor(self):
    self.assertEquals("", ground(0))

  def test_next_floor_for_1st(self):
    self.assertEquals(" ***", next_floor("*****"))

  def test_next_floor_for_2nd(self):
    self.assertEquals("  *", next_floor(" ***"))

def _4_floor_pyramid():
  return \
    "   *\n" + \
    "  ***\n" + \
    " *****\n" + \
    "*******"

def _3_floor_pyramid():
  return \
    "  *\n" + \
    " ***\n" + \
    "*****"\

if __name__ == '__main__':
  unittest.main()