Convert collocationScoreQuery list arguments to R vectors
Change-Id: Ibd6a0b23787b8d2568835f05c829a7f03c2825e3
diff --git a/KorAPClient/__init__.py b/KorAPClient/__init__.py
index 8bb5f5f..8b2b6b5 100644
--- a/KorAPClient/__init__.py
+++ b/KorAPClient/__init__.py
@@ -90,7 +90,7 @@
with localconverter(robjects.default_converter + pandas2ri.converter):
return robjects.conversion.rpy2py(KorAPClient.frequencyQuery(self, *args, **kwargs))
- def collocationScoreQuery(self, *args, **kwargs):
+ def collocationScoreQuery(self, node, collocate, vc="", **kwargs):
"""Get collocation scores for given node(s) and collocate(s).
- **node** - target word
@@ -114,7 +114,14 @@
```
"""
with localconverter(robjects.default_converter + pandas2ri.converter):
- return robjects.conversion.rpy2py(KorAPClient.collocationScoreQuery(self, *args, **kwargs))
+ if type(node) is list:
+ node = robjects.StrVector(node)
+ if type(collocate) is list:
+ collocate = robjects.StrVector(collocate)
+ if type(vc) is list:
+ vc = robjects.StrVector(vc)
+
+ return robjects.conversion.rpy2py(KorAPClient.collocationScoreQuery(self, node, collocate, vc, **kwargs))
def corpusQuery(self, *args, **kwargs):
"""Query search term(s).
diff --git a/KorAPClient/tests/test_korapclient.py b/KorAPClient/tests/test_korapclient.py
index 1b9774e..b359a19 100644
--- a/KorAPClient/tests/test_korapclient.py
+++ b/KorAPClient/tests/test_korapclient.py
@@ -1,7 +1,6 @@
import unittest
from KorAPClient import KorAPConnection
-
class TestKorAPClient(unittest.TestCase):
def setUp(self):
self.kcon = KorAPConnection(verbose=True)
@@ -22,9 +21,10 @@
self.assertGreater(df['pmi'][0], 10)
self.assertLess(df['pmi'][0], 20)
- @unittest.skip("Conversion error in rpy2py")
def test_collocation_score_query_multi_collocates(self):
df = self.kcon.collocationScoreQuery("Ameisenplage", ["einer", "heimgesucht"], leftContextSize=1, rightContextSize=1)
+ self.assertEqual(df['collocate'][1], 'heimgesucht')
+ self.assertGreater(df['pmi'][1], df['pmi'][0])
def test_corpus_stats(self):
df = self.kcon.corpusStats(**{"as.df": True})
@@ -35,6 +35,9 @@
ch_tokens = self.kcon.corpusStats(vc='pubPlaceKey="CH"', **{"as.df": True})['tokens'][0]
self.assertGreater(de_tokens, ch_tokens)
+ def test_corpus_stats_with_vc(self):
+ tokens = self.kcon.corpusStats(vc=['pubPlaceKey="DE"', 'pubPlaceKey="CH"'], **{"as.df": True})['tokens']
+ self.assertGreater(tokens[0], tokens[1])
if __name__ == '__main__':
unittest.main()