diff --git a/nemubot/tools/human.py b/nemubot/tools/human.py index b751409..588ac1f 100644 --- a/nemubot/tools/human.py +++ b/nemubot/tools/human.py @@ -57,3 +57,13 @@ def word_distance(str1, str2): ) return d[len(str1)][len(str2)] + + +def guess(pattern, expect): + if len(expect): + se = sorted([(e, word_distance(pattern, e)) for e in expect], key=lambda x: x[1]) + _, m = se[0] + for e, wd in se: + if wd > m or wd > 1 + len(pattern) / 4: + break + yield e diff --git a/nemubot/tools/test_human.py b/nemubot/tools/test_human.py index 06b48d8..8ebdd49 100644 --- a/nemubot/tools/test_human.py +++ b/nemubot/tools/test_human.py @@ -1,6 +1,6 @@ import unittest -from nemubot.tools.human import size, word_distance +from nemubot.tools.human import guess, size, word_distance class TestHuman(unittest.TestCase): @@ -31,6 +31,10 @@ class TestHuman(unittest.TestCase): self.assertEqual(word_distance("long", "short"), 4) self.assertEqual(word_distance("long", "short"), word_distance("short", "long")) + def test_guess(self): + self.assertListEqual([g for g in guess("drunk", ["eat", "drink"])], ["drink"]) + self.assertListEqual([g for g in guess("drunk", ["long", "short"])], []) + if __name__ == '__main__': unittest.main()