2016-07-29

Pythonでの演算子のオーバーロード

Python3.x

Pythonにも演算子のオーバーロードあるって知ったので、2次元ベクトルクラス作ってみた
Python3系で書いたけど2系でもいけるよね？

# -*- coding: utf-8 -*-

class Vec2D:
	def __init__(self, x, y):
		self.x = x
		self.y = y
	
	def __str__(self):
		return str(list((self.x, self.y)))
		
	def __eq__(self, other):
		return self.x == other.x and self.y == other.y
	def __ne__(self, other):
		return self.x != other.x or self.y != other.y
		
	def __add__(self, other):
		return Vec2D(self.x + other.x, self.x + other.y)
	def __iadd__(self, other):
		self.x += other.x
		self.y += other.y
		return self
		
	def __sub__(self, other):
		return Vec2D(self.x - other.x, self.y - other.y)
	def __isub__(self, other):
		self.x -= other.x
		self.y -= other.y
		return self
		
	def __mul__(self, other):
		return Vec2D(self.x * other, self.y * other)
	def __imul__(self, other):
		self.x *= other
		self.y *= other
		return self
		
	def __div__(self, other):
		return Vec2D(self.x / other, self.y / other)
	def __truediv__(self, other):
		return Vec2D(self.x / other, self.y / other)
	def __idiv__(self, other):
		self.x /= other
		self.y /= other
		return self

p1 = Vec2D(1, 2)
p2 = Vec2D(3, 4)
p3 = Vec2D(1, 2)

print("p1 :", p1)
print("p2 :", p2)
print("p3 :", p3)

print("p1 == p2 :", p1 == p2)
print("p1 == p3 :", p1 == p3)
print("p1 is p2 :", p1 is p2)
print("p1 is p3 :", p1 is p3)

print("p1 + p2 :", p1 + p2)
print("p1 - p2 :", p1 - p2)
p1 += p2
print("p1 += p2(p1) :", p1)
p1 -= p2
print("p1 -= p2(p1) :", p1)

print("p1 * 2 :", p1 * 2)
print("p1 / 2 :", p1 / 2)
p1 *= 2
print("p1 *= 2(p1) :", p1)
p1 /= 2
print("p1 /= 2(p1) :", p1)

下半分がテストコード
基本的なところしか実装してないけどまあこれでも使えるでしょ
演算結果はリストで吐き出しています

割り算は __div__() じゃなくて __truediv__() を定義しないとだめでした
文字の途中のハイライトわかんないのでそのままで許して

実行結果はこんな感じ

p1 : [1, 2]
p2 : [3, 4]
p3 : [1, 2]
p1 == p2 : False
p1 == p3 : True
p1 is p2 : False
p1 is p3 : False
p1 + p2 : [4, 5]
p1 - p2 : [-2, -2]
p1 += p2(p1) : [4, 6]
p1 -= p2(p1) : [1, 2]
p1 * 2 : [2, 4]
p1 / 2 : [0.5, 1.0]
p1 *= 2(p1) : [2, 4]
p1 /= 2(p1) : [1.0, 2.0]