Google Colaboratoryでコード整形

Google ColaboratoryやJupyter Notebook上で書いたコードをその場できれいに整形するマジックコマンドを作りました。

一言で説明すると…

こんな感じのいろいろ終了してるコードが…

x = {  'a':37,'b':42,

'c':927}

y = 'hello ''world'
z = 'hello '+'world'
a = 'hello {}'.format('world')
class foo  (     object  ):
 def f    (self   ):
   return       37*-+2
 def g(self, x,y=42):
     return y
def f  (   a ) :
 return      37+-+a[42-x :  y**3]

なんということでしょう！！（きれいになります）

x = {'a': 37, 'b': 42, 'c': 927}

y = 'hello ' 'world'
z = 'hello ' + 'world'
a = 'hello {}'.format('world')


class foo(object):
   def f(self):
       return 37 * -+2

   def g(self, x, y=42):
       return y


def f(a):
   return 37 + -+a[42 - x:y**3]

autopep8など、これまでの整形ツールはJupyter Notebookではワンタッチで使えましたが、Google Colabではうまく使えませんでした。。。（いい方法をご存じの方は教えてください）
そこで、マジックコマンド（%%timeitみたいなやつ）作りました。

使い方

まず最初のセルで以下のコマンドを実行。

!pip install yapf

次のセルで以下のコードをコピペして実行。

from IPython.core.magic import register_cell_magic
from yapf.yapflib.yapf_api import FormatCode


@register_cell_magic
def fmt(line, cell):
   """
   My formatter cell magic comannd.
   Please install yapf before using this magic command.
   !pip install yapf
   """
   print(FormatCode(cell, style_config='pep8')[0])

すると次のセル以降では「%%fmt」と最初の行に入力するだけで、そのセルのPythonコードを整形して出力してくれます。

あとは出力をコピペして実行するだけ。

詳細

このマジックコマンドでは、yapfという整形ツールを使用しています。
style_config='pep8'の部分を変えることで整形に用いるコーディングスタイルを変更できます。

なお、冒頭のコード整形の例で用いた終了しているソースコードは、 https://github.com/google/yapf#example のものを使用しています。

おまけ①

yapfでは弱いimport文の整形を行うisortと、あとノリでautopep8も使って3段階の整形をするバージョンです。「%%max_fmt」で使用できます。こちらの場合、「!pip install yapf isort autopep8」が必要です。

import autopep8
import isort
from IPython.core.magic import register_cell_magic
from yapf.yapflib.yapf_api import FormatCode


@register_cell_magic
def max_fmt(line, cell):
   """
   My formatter cell magic comannd.
   Please install autopep8, isort and yapf before using this magic command.
   !pip install autopep8 isort yapf
   """
   ret = isort.code(cell)
   ret = autopep8.fix_code(ret)
   print(FormatCode(ret, style_config='pep8')[0])

おまけ②

以下のように、一旦保存して整形することもできます。使い方は同じです。
具体的に言うと、セルに書かれているコードを一旦"myfomatter.fmt"というファイルに保存してyapfに引数で与えることで実現しています。
なお、この一時ファイル"myfomatter.fmt"は、万が一同一のファイルがあっても上書きしないように"os.path.exists(file_name)"で存在確認をしています。
ここでは、使用後の一時ファイルは即座に削除されるようにしています。
（あまり使い道はないかと思いますが）シェルのコマンドでいろいろしたいときは便利かもしれません。

import os

from IPython.core.magic import register_cell_magic


@register_cell_magic
def fmt_tmp(line, cell):
   """
   My formatter cell magic comannd.
   Please install yapf before using this magic command.
   !pip install yapf
   """
   original_file_name = "myformatter.fmt"
   file_name = original_file_name
   for i in range(10000):
       if os.path.exists(file_name):
           name, ext = os.path.splitext(original_file_name)
           name += str(i)
           file_name = name + ext
           continue
       else:
           break
   else:
       return f"Please rename {file_name}"
   with open(file_name, "w") as f:
       f.write(cell)
   !yapf $file_name
   !rm $file_name