YT-21446: Support py udfs generator is os ydb

ab1a950f25a9a5f4a9ca38655f443ff23cece827

1 files changed, 19 insertions, 6 deletions

diff --git a/build/plugins/yql_python_udf.py b/build/plugins/yql_python_udf.py
index 50f5526e89..18c2e9edce 100644
--- a/build/plugins/yql_python_udf.py
+++ b/build/plugins/yql_python_udf.py
@@ -1,4 +1,5 @@
 from _common import sort_by_keywords
+import ymake
 
 
 def get_or_default(kv, name, default):
@@ -17,9 +18,19 @@ def onregister_yql_python_udf(unit, *args):
     use_arcadia_python = unit.get('USE_ARCADIA_PYTHON') == 'yes'
     py3 = unit.get('PYTHON3') == 'yes'
 
+    if unit.get('OPENSOURCE'):
+        if add_libra_modules:
+            ymake.report_configure_error('Libra modules are not supported in opensource python UDFs')
+            add_libra_modules = False
+
+    yql_base_dir = unit.get('YQL_BASE_DIR')
+    yql_python_dir = unit.get('YQL_PYTHON_DIR')
+    if not yql_python_dir:
+        yql_python_dir = '/'.join([yql_base_dir, 'udfs/common/python'])
+
     unit.onyql_abi_version(['2', '27', '0'])
-    unit.onpeerdir(['yql/udfs/common/python/python_udf'])
-    unit.onpeerdir(['contrib/ydb/library/yql/public/udf'])
+    unit.onpeerdir(['/'.join([yql_python_dir, '/python_udf'])])
+    unit.onpeerdir(['/'.join([yql_base_dir, '/public/udf'])])
 
     if add_libra_modules:
         unit.onpeerdir(['quality/user_sessions/libra_arc/noyql'])
@@ -28,16 +39,16 @@ def onregister_yql_python_udf(unit, *args):
     if use_arcadia_python:
         flavor = 'Arcadia'
         unit.onpeerdir(
-            ['library/python/runtime', 'yql/udfs/common/python/main']
+            ['library/python/runtime', '/'.join([yql_python_dir, '/main'])]
             if not py3
-            else ['library/python/runtime_py3', 'yql/udfs/common/python/main_py3']
+            else ['library/python/runtime_py3', '/'.join([yql_python_dir, '/main_py3'])]
         )
     else:
         flavor = 'System'
 
     output_includes = [
-        'yql/udfs/common/python/python_udf/python_udf.h',
-        'contrib/ydb/library/yql/public/udf/udf_registrator.h',
+        '/'.join([yql_python_dir, '/python_udf/python_udf.h']),
+        '/'.join([yql_base_dir, '/public/udf/udf_registrator.h']),
     ]
     if add_libra_modules:
         output_includes.append('yql/udfs/quality/libra/module/module.h')
@@ -52,6 +63,8 @@ def onregister_yql_python_udf(unit, *args):
             resource_name,
             path,
             libra_flag,
+            yql_base_dir,
+            yql_python_dir,
             'OUT',
             path,
             'OUTPUT_INCLUDES',